高斯

GaussianPrior

基类：PriorDistribution

表示高斯先验分布的子类。

源代码位于 bionemo/moco/distributions/prior/continuous/gaussian.py

class GaussianPrior(PriorDistribution):
    """A subclass representing a Gaussian prior distribution."""

    def __init__(
        self,
        mean: Float = 0.0,
        std: Float = 1.0,
        center: Bool = False,
        rng_generator: Optional[torch.Generator] = None,
    ) -> None:
        """Gaussian prior distribution.

        Args:
            mean (Float): The mean of the Gaussian distribution. Defaults to 0.0.
            std (Float): The standard deviation of the Gaussian distribution. Defaults to 1.0.
            center (bool): Whether to center the samples around the mean. Defaults to False.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
        """
        self.mean = mean
        self.std = std
        self.center = center
        self.rng_generator = rng_generator

    def sample(
        self,
        shape: Tuple,
        mask: Optional[Tensor] = None,
        device: Union[str, torch.device] = "cpu",
        rng_generator: Optional[torch.Generator] = None,
    ) -> Tensor:
        """Generates a specified number of samples from the Gaussian prior distribution.

        Args:
            shape (Tuple): The shape of the samples to generate.
            device (str): cpu or gpu.
            mask (Optional[Tensor]): An optional mask to apply to the samples. Defaults to None.
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

        Returns:
            Float: A tensor of samples.
        """
        if rng_generator is None:
            rng_generator = self.rng_generator
        samples = torch.randn(*shape, device=device, generator=rng_generator)
        if self.std != 1:
            samples = samples * self.std
        if self.mean != 0:
            samples = samples + self.mean

        if self.center:
            samples = remove_center_of_mass(samples, mask)
        if mask is not None:
            samples = samples * mask.unsqueeze(-1)
        return samples

__init__(mean=0.0, std=1.0, center=False, rng_generator=None)

高斯先验分布。

参数

名称	类型	描述	默认值
mean	浮点数	高斯分布的均值。默认为 0.0。	0.0
std	浮点数	高斯分布的标准差。默认为 1.0。	1.0
center	布尔值	是否将样本围绕均值居中。默认为 False。	False
rng_generator	Optional[Generator]	用于可重复采样的可选 :class:torch.Generator。默认为 None。	None

源代码位于 bionemo/moco/distributions/prior/continuous/gaussian.py

def __init__(
    self,
    mean: Float = 0.0,
    std: Float = 1.0,
    center: Bool = False,
    rng_generator: Optional[torch.Generator] = None,
) -> None:
    """Gaussian prior distribution.

    Args:
        mean (Float): The mean of the Gaussian distribution. Defaults to 0.0.
        std (Float): The standard deviation of the Gaussian distribution. Defaults to 1.0.
        center (bool): Whether to center the samples around the mean. Defaults to False.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.
    """
    self.mean = mean
    self.std = std
    self.center = center
    self.rng_generator = rng_generator

sample(shape, mask=None, device='cpu', rng_generator=None)

从高斯先验分布生成指定数量的样本。

参数

名称	类型	描述	默认值
shape	元组	要生成的样本的形状。	必需
device	字符串	cpu 或 gpu。	'cpu'
mask	Optional[Tensor]	应用于样本的可选掩码。默认为 None。	None
rng_generator	Optional[Generator]	用于可重复采样的可选 :class:torch.Generator。默认为 None。	None

返回值

名称	类型	描述
浮点数	张量	样本的张量。

源代码位于 bionemo/moco/distributions/prior/continuous/gaussian.py

def sample(
    self,
    shape: Tuple,
    mask: Optional[Tensor] = None,
    device: Union[str, torch.device] = "cpu",
    rng_generator: Optional[torch.Generator] = None,
) -> Tensor:
    """Generates a specified number of samples from the Gaussian prior distribution.

    Args:
        shape (Tuple): The shape of the samples to generate.
        device (str): cpu or gpu.
        mask (Optional[Tensor]): An optional mask to apply to the samples. Defaults to None.
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

    Returns:
        Float: A tensor of samples.
    """
    if rng_generator is None:
        rng_generator = self.rng_generator
    samples = torch.randn(*shape, device=device, generator=rng_generator)
    if self.std != 1:
        samples = samples * self.std
    if self.mean != 0:
        samples = samples + self.mean

    if self.center:
        samples = remove_center_of_mass(samples, mask)
    if mask is not None:
        samples = samples * mask.unsqueeze(-1)
    return samples