pytorch中torch.bernoulli的使用方法-随机丢弃数据中的部分数据

今天在学习深度学习代码时，使用到如何将数据中的嵌入信息进行随机的丢弃。
在 PyTorch 中，torch.bernoulli 是用于生成服从伯努利分布的随机张量的函数。伯努利分布是一种二元分布，产生的值为0或1，具有概率 p 和 1-p。

语法

torch.bernoulli(input, *, generator=None, out=None) -> Tensor

input（张量）：一个包含概率值的张量。每个元素表示生成的随机数为1的概率。张量的元素值必须在区间 [0, 1] 中。生成的随机数在对应位置上的概率为该位置上的值。

generator（Generator，可选）：用于生成随机数的随机数生成器。

out（张量，可选）：输出张量。如果提供了此参数，则结果将被写入这个张量。否则，将创建一个新的张量并返回。

返回值

一个与输入张量具有相同形状的张量，其中每个元素都是根据输入张量中对应位置上的概率值生成的随机数。生成的随机数为1的概率由输入张量中对应位置上的值确定。

具体应用代码

def dropout(x, drop_prob, shared_axes=[], training=False):
    """
    Apply dropout to input tensor.
    Parameters
    ----------
    input_tensor: ``torch.FloatTensor``
        A tensor of shape ``(batch_size, ..., num_timesteps, embedding_dim)``
    Returns
    -------
    output: ``torch.FloatTensor``
        A tensor of shape ``(batch_size, ..., num_timesteps, embedding_dim)`` with dropout applied.
    """
    if drop_prob == 0 or drop_prob == None or (not training):
        return x
    sz = list(x.size())
    for i in shared_axes:
        sz[i] = 1
    x_new = x.new(*sz)
    x_bernoulli = x_new.bernoulli_(1. - drop_prob)
    mask = x_bernoulli.div_(1. - drop_prob)
    mask = mask.expand_as(x)
    return x * mask