TorchDiff 2.0.0__py3-none-any.whl

sde/hyper_param.py

@@ -0,0 +1,200 @@
+import torch
+import torch.nn as nn
+
+
+
+
+class HyperParamsSDE(nn.Module):
+    """Hyperparameters for SDE-based generative models.
+
+    Manages the noise schedule and SDE-specific parameters for score-based generative
+    models, including beta and sigma schedules, time steps, and variance computations,
+    as described in Song et al. (2021). Supports trainable or fixed beta schedules and
+    multiple scheduling methods for flexible noise control.
+
+    Parameters
+    ----------
+    num_steps : int, optional
+        Number of diffusion steps (default: 1000).
+    beta_start : float, optional
+        Starting value for beta schedule (default: 1e-4).
+    beta_end : float, optional
+        Ending value for beta schedule (default: 0.02).
+    trainable_beta : bool, optional
+        Whether the beta schedule is trainable (default: False).
+    beta_method : str, optional
+        Method for computing the beta schedule (default: "linear").
+        Supported methods: "linear", "sigmoid", "quadratic", "constant", "inverse_time".
+    sigma_start : float, optional
+        Starting value for sigma schedule for VE method (default: 1e-3).
+    sigma_end : float, optional
+        Ending value for sigma schedule for VE method (default: 10.0).
+    start : float, optional
+        Start of the time interval for SDE integration (default: 0.0).
+    end : float, optional
+        End of the time interval for SDE integration (default: 1.0).
+
+    Attributes
+    ----------
+    num_steps : int
+        Number of diffusion steps.
+    beta_start : float
+        Minimum beta value.
+    beta_end : float
+        Maximum beta value.
+    trainable_beta : bool
+        Whether the beta schedule is trainable.
+    beta_method : str
+        Method used for beta schedule computation.
+    sigma_start : float
+        Minimum sigma value for VE method.
+    sigma_end : float
+        Maximum sigma value for VE method.
+    start : float
+        Start of the time interval.
+    end : float
+        End of the time interval.
+    betas : torch.Tensor
+        Beta schedule values, shape (num_steps,). Trainable if `trainable_beta` is True,
+        otherwise a fixed buffer.
+    cum_betas : torch.Tensor, optional
+        Cumulative sum of betas scaled by `dt`, shape (num_steps,). Available if
+        `trainable_beta` is False.
+    sigmas : torch.Tensor, optional
+        Sigma schedule for VE method, shape (num_steps,). Available if
+        `trainable_beta` is False.
+    time : torch.Tensor
+        Time points for SDE integration, shape (num_steps,).
+    dt : float
+        Time step size for SDE integration, computed as (end - start) / num_steps.
+
+    Raises
+    ------
+    ValueError
+        If `beta_start` or `beta_end` do not satisfy 0 < beta_start < beta_end,
+        `sigma_start` or `sigma_end` do not satisfy 0 < sigma_start < sigma_end,
+        or `num_steps` is not positive.
+    """
+    def __init__(self, num_steps=1000, beta_start=1e-4, beta_end=0.02, trainable_beta=False, beta_method="linear",
+                 sigma_start=1e-3, sigma_end=10.0, start=0.0, end=1.0):
+        super().__init__()
+        self.num_steps = num_steps
+        self.beta_start = beta_start
+        self.beta_end = beta_end
+        self.trainable_beta = trainable_beta
+        self.beta_method = beta_method
+        self.sigma_start = sigma_start
+        self.sigma_end = sigma_end
+        self.start = start
+        self.end = end
+
+        if not (0 < self.beta_start < self.beta_end):
+            raise ValueError(f"beta_start ({self.beta_start}) and beta_end ({self.beta_end}) must satisfy 0 < start < end")
+        if not (0 < self.sigma_start < self.sigma_end):
+            raise ValueError(f"sigma_start ({self.sigma_start}) and sigma_end ({self.sigma_end}) must satisfy 0 < start < end")
+        if self.num_steps <= 0:
+            raise ValueError(f"num_steps ({self.num_steps}) must be positive")
+
+        beta_range = (beta_start, beta_end)
+        betas_init = self.compute_beta_schedule(beta_range, num_steps, beta_method)
+        self.time = torch.linspace(self.start, self.end, self.num_steps, dtype=torch.float32)
+        self.dt = (self.end - self.start) / self.num_steps
+
+        if trainable_beta:
+            self.betas = nn.Parameter(betas_init)
+        else:
+            self.register_buffer('betas', betas_init)
+            self.register_buffer('cum_betas', torch.cumsum(betas_init, dim=0) * self.dt)
+            self.register_buffer("sigmas", self.sigma_start * (self.sigma_end / self.sigma_start) ** self.time)
+
+    def compute_beta_schedule(self, beta_range, num_steps, method):
+        """Computes the beta schedule based on the specified method.
+
+        Generates a sequence of beta values for the SDE noise schedule using the chosen
+        method, ensuring values are clamped within the specified range.
+
+        Parameters
+        ----------
+        beta_range : tuple
+            Tuple of (min_beta, max_beta) specifying the valid range for beta values.
+        num_steps : int
+            Number of diffusion steps.
+        method : str
+            Method for computing the beta schedule. Supported methods:
+            "linear", "sigmoid", "quadratic", "constant", "inverse_time".
+
+        Returns
+        -------
+        torch.Tensor
+            Tensor of beta values, shape (num_steps,).
+
+        Raises
+        ------
+        ValueError
+            If `method` is not one of the supported beta schedule methods.
+        """
+        beta_min, beta_max = beta_range
+        if method == "sigmoid":
+            x = torch.linspace(-6, 6, num_steps)
+            beta = torch.sigmoid(x) * (beta_max - beta_min) + beta_min
+        elif method == "quadratic":
+            x = torch.linspace(beta_min ** 0.5, beta_max ** 0.5, num_steps)
+            beta = x ** 2
+        elif method == "constant":
+            beta = torch.full((num_steps,), beta_max)
+        elif method == "inverse_time":
+            beta = 1.0 / torch.linspace(num_steps, 1, num_steps)
+            beta = beta_min + (beta_max - beta_min) * (beta - beta.min()) / (beta.max() - beta.min())
+        elif method == "linear":
+            beta = torch.linspace(beta_min, beta_max, num_steps)
+        else:
+            raise ValueError(f"Unknown beta_method: {method}. Supported: linear, sigmoid, quadratic, constant, inverse_time")
+        beta = torch.clamp(beta, min=beta_min, max=beta_max)
+        return beta
+
+    def constrain_betas(self):
+        """Constrains trainable betas to a valid range during training.
+
+        Ensures that trainable beta values remain within the specified range
+        [beta_start, beta_end] by clamping them in-place.
+
+        Notes
+        -----
+        This method only applies when `trainable_beta` is True.
+        """
+        if self.trainable_beta:
+            with torch.no_grad():
+                self.betas.clamp_(min=self.beta_start, max=self.beta_end)
+
+    def get_variance(self, time_steps, method):
+        """Computes the variance for the specified SDE method at given time steps.
+
+        Calculates the variance used in SDE diffusion processes based on the method
+        (VE, VP, or sub-VP), leveraging the sigma or cumulative beta schedules.
+
+        Parameters
+        ----------
+        time_steps : torch.Tensor
+            Tensor of time step indices (long), shape (batch_size,), where each value
+            is in the range [0, num_steps - 1].
+        method : str
+            SDE method to compute variance for. Supported methods: "ve", "vp", "sub-vp".
+
+        Returns
+        -------
+        torch.Tensor
+            Variance values for the specified time steps, shape (batch_size,).
+
+        Raises
+        ------
+        ValueError
+            If `method` is not one of the supported methods ("ve", "vp", "sub-vp").
+        """
+        if method == "ve":
+            return self.sigmas[time_steps] ** 2
+        elif method == "vp":
+            return 1 - torch.exp(-self.cum_betas[time_steps])
+        elif method == "sub-vp":
+            return 1 - torch.exp(-2 * self.cum_betas[time_steps])
+        else:
+            raise ValueError(f"Unknown method: {method}")