torch.bernoulli

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

生成伯努利分布中的二进制随机数（0或1）。

张量 input 应包含用于生成二进制随机数的概率值。因此，input 中的所有值必须在 0 到 1 的范围内：$0 \leq \text{input}_i \leq 1$。

输出张量的第 \(\text{i}^{\text{th}}\) 元素将根据 input 中给出的第 \(\text{i}^{\text{th}}\) 概率值抽取一个$1$值。

$\text{out}_{i} \sim \mathrm{Bernoulli}(p = \text{input}_{i})$

返回的 out 张量仅包含 0 和 1，其形状与 input 相同。

out 可以是整数类型，但 input 必须是浮点类型。

参数

输入 (Tensor) – 表示伯努利分布概率值的输入张量

关键字参数

• generator (torch.Generator, 可选) – 用于样本采集的伪随机数生成器

• out (Tensor, 可选) – 指定输出张量。

示例：

>>> a = torch.empty(3, 3).uniform_(0, 1)  # generate a uniform random matrix with range [0, 1]
>>> a
tensor([[ 0.1737,  0.0950,  0.3609],
        [ 0.7148,  0.0289,  0.2676],
        [ 0.9456,  0.8937,  0.7202]])
>>> torch.bernoulli(a)
tensor([[ 1.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 1.,  1.,  1.]])

>>> a = torch.ones(3, 3) # probability of drawing "1" is 1
>>> torch.bernoulli(a)
tensor([[ 1.,  1.,  1.],
        [ 1.,  1.,  1.],
        [ 1.,  1.,  1.]])
>>> a = torch.zeros(3, 3) # probability of drawing "1" is 0
>>> torch.bernoulli(a)
tensor([[ 0.,  0.,  0.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]])