自定义

DiscreteCustomPrior

基类：DiscretePriorDistribution

一个子类，表示离散自定义先验分布。

此类允许创建具有自定义概率质量函数的先验分布，该函数由 prior_dist 张量定义。例如，如果我的数据有 4 个类别，并且我想要 [.3, .2, .4, .1] 作为 4 个类别的概率。

源代码在 bionemo/moco/distributions/prior/discrete/custom.py 中

class DiscreteCustomPrior(DiscretePriorDistribution):
    """A subclass representing a discrete custom prior distribution.

    This class allows for the creation of a prior distribution with a custom
    probability mass function defined by the `prior_dist` tensor. For example if my data has 4 classes and I want [.3, .2, .4, .1] as the probabilities of the 4 classes.
    """

    def __init__(self, prior_dist: Tensor, num_classes: int = 10) -> None:
        """Initializes a DiscreteCustomPrior distribution.

        Args:
            prior_dist: A tensor representing the probability mass function of the prior distribution.
            num_classes: The number of classes in the prior distribution. Defaults to 10.

        Note:
            The `prior_dist` tensor should have a sum close to 1.0, as it represents a probability mass function.
        """
        super().__init__(num_classes, prior_dist)
        if torch.sum(self.prior_dist).item() - 1.0 > 1e-5:
            raise ValueError("Prior distribution probabilities do not sum up to 1.0")

    def sample(
        self,
        shape: Tuple,
        mask: Optional[Tensor] = None,
        device: Union[str, torch.device] = "cpu",
        rng_generator: Optional[torch.Generator] = None,
    ) -> Tensor:
        """Samples from the discrete custom prior distribution.

        Args:
            shape: A tuple specifying the shape of the samples to generate.
            mask: An optional tensor mask to apply to the samples, broadcastable to the sample shape. Defaults to None.
            device: The device on which to generate the samples, specified as a string or a :class:`torch.device`. Defaults to "cpu".
            rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

        Returns:
            A tensor of samples drawn from the prior distribution.
        """
        samples = (
            torch.multinomial(self.prior_dist, math.prod(shape), replacement=True, generator=rng_generator)
            .to(device)
            .reshape(shape)
        )
        if mask is not None:
            samples = samples * mask[(...,) + (None,) * (len(samples.shape) - len(mask.shape))]
        return samples

__init__(prior_dist, num_classes=10)

初始化 DiscreteCustomPrior 分布。

参数

名称	类型	描述	默认值
prior_dist	张量	表示先验分布的概率质量函数的张量。	必需
num_classes	整数	先验分布中的类别数量。默认为 10。	10

注意

prior_dist 张量的总和应接近 1.0，因为它表示概率质量函数。

源代码在 bionemo/moco/distributions/prior/discrete/custom.py 中

def __init__(self, prior_dist: Tensor, num_classes: int = 10) -> None:
    """Initializes a DiscreteCustomPrior distribution.

    Args:
        prior_dist: A tensor representing the probability mass function of the prior distribution.
        num_classes: The number of classes in the prior distribution. Defaults to 10.

    Note:
        The `prior_dist` tensor should have a sum close to 1.0, as it represents a probability mass function.
    """
    super().__init__(num_classes, prior_dist)
    if torch.sum(self.prior_dist).item() - 1.0 > 1e-5:
        raise ValueError("Prior distribution probabilities do not sum up to 1.0")

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

从离散自定义先验分布中采样。

参数

名称	类型	描述	默认值
shape	元组	指定要生成的样本形状的元组。	必需
mask	Optional[张量]	可选的张量掩码，应用于样本，可广播到样本形状。默认为 None。	无
device	Union[str, device]	在其上生成样本的设备，指定为字符串或 :class:torch.device。默认为 "cpu"。	'cpu'
rng_generator	Optional[Generator]	可选的 :class:torch.Generator，用于可重复的采样。默认为 None。	无

返回值

类型	描述
张量	从先验分布中抽取的样本张量。

源代码在 bionemo/moco/distributions/prior/discrete/custom.py 中

def sample(
    self,
    shape: Tuple,
    mask: Optional[Tensor] = None,
    device: Union[str, torch.device] = "cpu",
    rng_generator: Optional[torch.Generator] = None,
) -> Tensor:
    """Samples from the discrete custom prior distribution.

    Args:
        shape: A tuple specifying the shape of the samples to generate.
        mask: An optional tensor mask to apply to the samples, broadcastable to the sample shape. Defaults to None.
        device: The device on which to generate the samples, specified as a string or a :class:`torch.device`. Defaults to "cpu".
        rng_generator: An optional :class:`torch.Generator` for reproducible sampling. Defaults to None.

    Returns:
        A tensor of samples drawn from the prior distribution.
    """
    samples = (
        torch.multinomial(self.prior_dist, math.prod(shape), replacement=True, generator=rng_generator)
        .to(device)
        .reshape(shape)
    )
    if mask is not None:
        samples = samples * mask[(...,) + (None,) * (len(samples.shape) - len(mask.shape))]
    return samples