PyTorch Batch 屏蔽选择实现

Question

我怎样才能执行 批处理 masked_select？

给定：

x = torch.tensor([[1., 2., 2., 2., 3.],
                  [1., 2., 4., 3., 2.]])

期望的输出是：

tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])

这是一种可能的方法：

x = torch.tensor([[1., 2., 2., 2., 3.],
                  [1., 2., 4., 3., 2.]])

ones = torch.tensor([[1., 1., 1., 1., 1.],
                     [1., 1., 1., 1., 1.]])

masks = torch.tensor([[ True, False, False, False,  True],
                      [ True, False,  True,  True, False]])

for i in range(x.size(0)):
    mask = masks[i]
    s = torch.masked_select(x[i], mask)
    ones[i][:s.size(0)] = s

有其他解决方案吗？

Answer 1

这类问题的主要问题是中间结果是非同质的：在您的批次中，元素将具有不同数量的掩码值。如果我们想应用 PyTorch 内置函数，这是一个问题。在这里，我提供了两种解决方案来执行此操作。

---

1- 使用list comprehension

通过适当数量的批处理元素，掩码和填充：

>>> pad = lambda v: F.pad(v, [0, len(m)-len(v)], value=1)
>>> torch.stack([pad(r[m]) for r, m in zip(x, masks)])
tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])

这相当简单，与您的方法相似。

---

2- 使用torch.scatter

向量化的替代方法是构造正确的值和索引张量，以便我们可以应用torch.scatter 并获得所需的结果。这里的诀窍是使用展平张量。从x 和masks 我们首先要访问nz 和idx，定义为：

• nz：来自x 的非屏蔽值（当然使用masks），即我们需要找到：

  tensor([1., 3., 1., 4., 3.]) 
  

• idx：它们在输出张量中的对应索引扁平化时。

  tensor([ 0,  1,  5,  6,  7])
  

然后我们可以使用 out = ones.scatter(dim=0, idx, nz) 之类的东西来应用分散，这将有效地执行：out[idx[i]] = nz[i]。

要构造nz，我们可以使用masks 直接用masks非零值索引索引x：

>>> nz = x[masks]
tensor([1., 3., 1., 4., 3.])

对于idx，它会有点棘手。我们可以对掩码本身进行排序、展平并使用torch.Tensor.nonzero 获取非零值。排序后，True 值会出现在每行的开头：

>>> idx = masks.sort(1, True).values.view(-1).nonzero()[:,0]
tensor([ 0,  1,  5,  6,  7])

最后我们可以申请torch.scatter并reshape得到想要的结果：

>>> torch.ones(x.numel()).scatter(0, idx, nz).view_as(x)
tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])

这里torch.scatter 的使用是有限的，因为输入是一维的。一种等效的方法是：

>>> o = torch.ones(x.numel())
>>> o[idx] = nz
>>> o.view_as(x)

完整的方法：

>>> idx = masks.sort(1, True)[0].view(-1).nonzero()[:,0]
>>> torch.ones(x.numel()).scatter(0, idx, x[masks]).view_as(x)
tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])