Функция bernoulli

Функция bernoulli генерирует тензор с элементами, принимающими значения 0 или 1 в соответствии с распределением Бернулли. Первым параметром передается тензор вероятностей (значения от 0 до 1), где каждый элемент определяет вероятность получения единицы. Вторым параметром можно указать генератор случайных чисел.

Синтаксис

torch.bernoulli(input, [generator])

Пример

Сгенерируем тензор на основе вероятностей:

import torch

torch.manual_seed(0)
probs = torch.tensor([0.1, 0.5, 0.9])
res = torch.bernoulli(probs)
print(res)

Результат выполнения кода:

tensor([0., 1., 1.])

Пример

Сгенерируем двумерный тензор:

import torch

torch.manual_seed(1)
probs = torch.tensor([
	[0.2, 0.8],
	[0.6, 0.4],
])
res = torch.bernoulli(probs)
print(res)

Результат выполнения кода:

tensor([
	[0., 1.],
	[1., 0.],
])

Пример

Использование функции с генератором случайных чисел:

import torch

torch.manual_seed(2)
generator = torch.Generator()
generator.manual_seed(42)
probs = torch.tensor([0.3, 0.7])
res = torch.bernoulli(probs, generator=generator)
print(res)

Результат выполнения кода:

tensor([0., 1.])

Пример

Получение бинарной маски на основе вероятностей:

import torch

torch.manual_seed(3)
probs = torch.tensor([0.0, 0.5, 1.0])
res = torch.bernoulli(probs)
print(res)

Результат выполнения кода:

tensor([0., 1., 1.])

Смотрите также

функцию rand,
которая генерирует тензор с равномерным распределением
функцию randn,
которая генерирует тензор с нормальным распределением
функцию randint,
которая генерирует тензор с целочисленными значениями
функцию multinomial,
которая генерирует выборку из мультиномиального распределения