I have two same-shaped PyTorch tensors A and B, and I'd like to create a same-shape "randomly mixed" tensor C where C[i,...] = A[i,...] with probability alpha or B[i,...] with probability 1-alpha. Is there some Pythonic way to do this compactly?
1 Answers
consider using torch.bernoulli to create a mask tensor:
import torch
prob = 0.8
x = torch.full((2, 6, 3), 10.2, dtype=torch.float)
y = torch.full((2, 6, 3), -1.6, dtype=torch.float)
mask = torch.bernoulli(torch.full(x.shape, prob)).int()
reverse_mask = torch.ones(x.shape).int() - mask
result = x * mask + y * reverse_mask
result is now:
[[[10.2000, 10.2000, 10.2000],
[10.2000, -1.6000, 10.2000],
[10.2000, 10.2000, -1.6000],
[-1.6000, 10.2000, -1.6000],
[10.2000, 10.2000, 10.2000],
[10.2000, 10.2000, 10.2000]],
[[10.2000, 10.2000, -1.6000],
[10.2000, 10.2000, 10.2000],
[10.2000, 10.2000, -1.6000],
[10.2000, -1.6000, 10.2000],
[-1.6000, 10.2000, 10.2000],
[10.2000, 10.2000, 10.2000]]]
Good Luck!
Related