从两个相同大小的张量中随机选择项目

Question

假设我们有两个大小相同的张量batch_size * 1。对于批量维度中的每个索引，我们希望在两个张量之间随机选择。我的解决方案是创建一个 indices 张量，其中包含随机的 0 或 1 大小的索引 batch_size 并将它们用于 index_select 从两个张量的串联中。然而，这样做我有cat张量的“观点”，并且解决方案最终变得非常“丑陋”：

import torch

bs = 8
a = torch.zeros(bs, 1)
print("a size", a.size())
b = torch.ones(bs, 1)

c = torch.cat([a, b], dim=-1)
print(c)
print("c size", c.size())

# create bs number of random 0 and 1's
indices = torch.randint(0, 2, [bs])
print("idxs size", indices.size())
print("idxs", indices)

# use `indices` to slice the `cat`ted tensor
d = c.view(1, -1).index_select(-1, indices).view(-1, 1)
print("d size", d.size())
print(d)

我想知道是否有更漂亮，更重要的是，更有效的解决方案。

Answer 1

发布我在PyTorch forums 得到的两个答案

import torch

bs = 8
a = torch.zeros(bs, 1)
b = torch.ones(bs, 1)
c = torch.cat([a, b], dim=-1)
choices_flat = c.view(-1)
# index = torch.randint(choices_flat.numel(), (bs,))
# or if replace = False
index = torch.randperm(choices_flat.numel())[:bs]
select = choices_flat[index]

print(select)

import torch

bs = 8
a = torch.zeros(bs, 1)
print("a size", a.size())
b = torch.ones(bs, 1)

idx = torch.randint(2 * bs, (bs,))

d = torch.cat([a, b])[idx] # [bs, 1]