离散噪声 schedules

DiscreteCosineNoiseSchedule

基类：DiscreteNoiseSchedule

余弦离散噪声 schedule。

源代码位于 bionemo/moco/schedules/noise/discrete_noise_schedules.py

class DiscreteCosineNoiseSchedule(DiscreteNoiseSchedule):
    """A cosine discrete noise schedule."""

    def __init__(self, nsteps: int, nu: Float = 1.0, s: Float = 0.008):
        """Initialize the CosineNoiseSchedule.

        Args:
            nsteps (int): Number of discrete steps.
            nu (Optional[Float]): Hyperparameter for the cosine schedule exponent (default is 1.0).
            s (Optional[Float]): Hyperparameter for the cosine schedule shift (default is 0.008).
        """
        super().__init__(nsteps=nsteps, direction=TimeDirection.DIFFUSION)
        self.nu = nu
        self.s = s

    def _generate_schedule(self, nsteps: Optional[int] = None, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Generate the cosine noise schedule.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        if nsteps is None:
            nsteps = self.nsteps
        steps = (
            nsteps + 1
        )  #! matches OpenAI code https://github.com/openai/improved-diffusion/blob/main/improved_diffusion/gaussian_diffusion.py#L62
        x = torch.linspace(0, nsteps, steps, device=device)
        alphas_cumprod = torch.cos(((x / nsteps) ** self.nu + self.s) / (1 + self.s) * torch.pi * 0.5) ** 2
        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
        betas = torch.clip(betas, 0.001, 0.999)
        return 1 - betas

    def _clip_noise_schedule(self, alphas2: Tensor, clip_value: Float = 0.001) -> Tensor:
        """For a noise schedule given by alpha^2, this clips alpha_t / alpha_t-1. This may help improve stability during sampling.

        Args:
            alphas2 (Tensor): The noise schedule given by alpha^2.
            clip_value (Optional[Float]): The minimum value for alpha_t / alpha_t-1 (default is 0.001).

        Returns:
            Tensor: The clipped noise schedule.
        """
        alphas2 = torch.cat([torch.ones(1, device=alphas2.device), alphas2], dim=0)

        alphas_step = alphas2[1:] / alphas2[:-1]

        alphas_step = torch.clamp(alphas_step, min=clip_value, max=1.0)
        alphas2 = torch.cumprod(alphas_step, dim=0)

        return alphas2

__init__(nsteps, nu=1.0, s=0.008)

初始化 CosineNoiseSchedule。

参数

名称	类型	描述	默认值
nsteps	int	离散步数。	必需
nu	Optional[Float]	余弦 schedule 指数的超参数（默认为 1.0）。	1.0
s	Optional[Float]	余弦 schedule 偏移的超参数（默认为 0.008）。	0.008

源代码位于 bionemo/moco/schedules/noise/discrete_noise_schedules.py

def __init__(self, nsteps: int, nu: Float = 1.0, s: Float = 0.008):
    """Initialize the CosineNoiseSchedule.

    Args:
        nsteps (int): Number of discrete steps.
        nu (Optional[Float]): Hyperparameter for the cosine schedule exponent (default is 1.0).
        s (Optional[Float]): Hyperparameter for the cosine schedule shift (default is 0.008).
    """
    super().__init__(nsteps=nsteps, direction=TimeDirection.DIFFUSION)
    self.nu = nu
    self.s = s

DiscreteLinearNoiseSchedule

基类：DiscreteNoiseSchedule

线性离散噪声 schedule。

源代码位于 bionemo/moco/schedules/noise/discrete_noise_schedules.py

class DiscreteLinearNoiseSchedule(DiscreteNoiseSchedule):
    """A linear discrete noise schedule."""

    def __init__(self, nsteps: int, beta_start: Float = 1e-4, beta_end: Float = 0.02):
        """Initialize the CosineNoiseSchedule.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            beta_start (Optional[int]): starting beta value. Defaults to 1e-4.
            beta_end (Optional[int]): end beta value. Defaults to 0.02.
        """
        super().__init__(nsteps=nsteps, direction=TimeDirection.DIFFUSION)
        self.beta_start = beta_start
        self.beta_end = beta_end

    def _generate_schedule(self, nsteps: Optional[int] = None, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Generate the cosine noise schedule.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        if nsteps is None:
            nsteps = self.nsteps
        betas = torch.linspace(self.beta_start, self.beta_end, nsteps, dtype=torch.float32, device=device)
        return 1 - betas

__init__(nsteps, beta_start=0.0001, beta_end=0.02)

初始化 CosineNoiseSchedule。

参数

名称	类型	描述	默认值
nsteps	Optional[int]	时间步数。如果为 None，则使用初始化时的值。	必需
beta_start	Optional[int]	起始 beta 值。默认为 1e-4。	0.0001
beta_end	Optional[int]	结束 beta 值。默认为 0.02。	0.02

源代码位于 bionemo/moco/schedules/noise/discrete_noise_schedules.py

def __init__(self, nsteps: int, beta_start: Float = 1e-4, beta_end: Float = 0.02):
    """Initialize the CosineNoiseSchedule.

    Args:
        nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
        beta_start (Optional[int]): starting beta value. Defaults to 1e-4.
        beta_end (Optional[int]): end beta value. Defaults to 0.02.
    """
    super().__init__(nsteps=nsteps, direction=TimeDirection.DIFFUSION)
    self.beta_start = beta_start
    self.beta_end = beta_end

DiscreteNoiseSchedule

基类：ABC

离散噪声 schedules 的基类。

源代码位于 bionemo/moco/schedules/noise/discrete_noise_schedules.py

class DiscreteNoiseSchedule(ABC):
    """A base class for discrete noise schedules."""

    def __init__(self, nsteps: int, direction: TimeDirection):
        """Initialize the DiscreteNoiseSchedule.

        Args:
           nsteps (int): number of discrete steps.
           direction (TimeDirection): required this defines in which direction the scheduler was built
        """
        self.nsteps = nsteps
        self.direction = string_to_enum(direction, TimeDirection)

    def generate_schedule(
        self,
        nsteps: Optional[int] = None,
        device: Union[str, torch.device] = "cpu",
        synchronize: Optional[TimeDirection] = None,
    ) -> Tensor:
        """Generate the noise schedule as a tensor.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
            synchronize (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction,
                this parameter allows to flip the direction to match the specified one (default is None).
        """
        schedule = self._generate_schedule(nsteps, device)
        if synchronize and self.direction != string_to_enum(synchronize, TimeDirection):
            return torch.flip(schedule, dims=[0])
        else:
            return schedule

    @abstractmethod
    def _generate_schedule(self, nsteps: Optional[int] = None, device: Union[str, torch.device] = "cpu") -> Tensor:
        """Generate the noise schedule tensor.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
        """
        pass

    def calculate_derivative(
        self,
        nsteps: Optional[int] = None,
        device: Union[str, torch.device] = "cpu",
        synchronize: Optional[TimeDirection] = None,
    ) -> Tensor:
        """Calculate the time derivative of the schedule.

        Args:
            nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
            device (Optional[str]): Device to place the schedule on (default is "cpu").
            synchronize (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction,
                this parameter allows to flip the direction to match the specified one (default is None).

        Returns:
            Tensor: A tensor representing the time derivative of the schedule.

        Raises:
            NotImplementedError: If the derivative calculation is not implemented for this schedule.
        """
        raise NotImplementedError("Derivative calculation is not implemented for this schedule.")

__init__(nsteps, direction)

初始化 DiscreteNoiseSchedule。

参数

名称	类型	描述	默认值
nsteps	int	离散步数。	必需
direction	TimeDirection	必需，这定义了构建 scheduler 的方向	必需

源代码位于 bionemo/moco/schedules/noise/discrete_noise_schedules.py

def __init__(self, nsteps: int, direction: TimeDirection):
    """Initialize the DiscreteNoiseSchedule.

    Args:
       nsteps (int): number of discrete steps.
       direction (TimeDirection): required this defines in which direction the scheduler was built
    """
    self.nsteps = nsteps
    self.direction = string_to_enum(direction, TimeDirection)

calculate_derivative(nsteps=None, device='cpu', synchronize=None)

计算 schedule 的时间导数。

参数

名称	类型	描述	默认值
nsteps	Optional[int]	时间步数。如果为 None，则使用初始化时的值。	None
device	Optional[str]	放置 schedule 的设备（默认为 "cpu"）。	'cpu'
synchronize	Optional[str]	用于同步 schedule 的 TimeDirection。如果 schedule 是用不同的方向定义的，则此参数允许翻转方向以匹配指定的方向（默认为 None）。	None

返回值

名称	类型	描述
Tensor	Tensor	表示 schedule 时间导数的张量。

Raises

类型	描述
NotImplementedError	如果未为此 schedule 实现导数计算。

源代码位于 bionemo/moco/schedules/noise/discrete_noise_schedules.py

def calculate_derivative(
    self,
    nsteps: Optional[int] = None,
    device: Union[str, torch.device] = "cpu",
    synchronize: Optional[TimeDirection] = None,
) -> Tensor:
    """Calculate the time derivative of the schedule.

    Args:
        nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").
        synchronize (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction,
            this parameter allows to flip the direction to match the specified one (default is None).

    Returns:
        Tensor: A tensor representing the time derivative of the schedule.

    Raises:
        NotImplementedError: If the derivative calculation is not implemented for this schedule.
    """
    raise NotImplementedError("Derivative calculation is not implemented for this schedule.")

generate_schedule(nsteps=None, device='cpu', synchronize=None)

生成噪声 schedule 作为张量。

参数

名称	类型	描述	默认值
nsteps	Optional[int]	时间步数。如果为 None，则使用初始化时的值。	None
device	Optional[str]	放置 schedule 的设备（默认为 "cpu"）。	'cpu'
synchronize	Optional[str]	用于同步 schedule 的 TimeDirection。如果 schedule 是用不同的方向定义的，则此参数允许翻转方向以匹配指定的方向（默认为 None）。	None

源代码位于 bionemo/moco/schedules/noise/discrete_noise_schedules.py

def generate_schedule(
    self,
    nsteps: Optional[int] = None,
    device: Union[str, torch.device] = "cpu",
    synchronize: Optional[TimeDirection] = None,
) -> Tensor:
    """Generate the noise schedule as a tensor.

    Args:
        nsteps (Optional[int]): Number of time steps. If None, uses the value from initialization.
        device (Optional[str]): Device to place the schedule on (default is "cpu").
        synchronize (Optional[str]): TimeDirection to synchronize the schedule with. If the schedule is defined with a different direction,
            this parameter allows to flip the direction to match the specified one (default is None).
    """
    schedule = self._generate_schedule(nsteps, device)
    if synchronize and self.direction != string_to_enum(synchronize, TimeDirection):
        return torch.flip(schedule, dims=[0])
    else:
        return schedule