作者:justmoon999 | 来源:互联网 | 2023-08-06 10:41
我有两个形状相同的 PyTorch 张量 A 和 B,我想创建一个形状相同的“随机混合”张量 C,其中 C[i,...] = A[i,...] 概率为 alpha或 B[i,...] 概率为 1-alpha。有没有一些 Pythonic 的方式来紧凑地做到这一点?
回答
考虑使用torch.bernoulli
来创建掩码张量:
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
就是现在:
[[[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]]]
祝你好运!