使用布尔张量对张量进行切片

Question

有两个张量 :inputs_tokens 是一批 20x300 的令牌 ID seq_A 是我的模型输出，大小为 [20, 300, 512]（批次中每个标记的 512 个向量）

seq_A.size()
Out[1]: torch.Size([20, 300, 512])

inputs_tokens.size()
torch.Size([20, 300])

我只想获取令牌 101 (CLS) 的向量，如下所示：

cls_tokens = (inputs_tokens == 101) 
cls_tokens
Out[4]: 
tensor([[ True, False, False,  ..., False, False, False],
       [ True, False, False,  ..., False, False, False],
       [ True, False, False,  ..., False, False, False], ...

如何对 seq_A 进行切片以仅获取每个批次的 cls_tokens 中为真的向量？ 当我这样做时

seq_A[cls_tokens].size()
Out[7]: torch.Size([278, 512]) 

但我仍然需要它以达到 [20 x N x 512] 的大小（否则我不知道它属于哪个样本）

Answer 1

TLDR；不能，所有序列必须沿给定轴具有相同的大小。

---

以这个简化的例子为例：

>>> inputs_tokens = torch.tensor([[  1, 101,  18, 101,   9],
                                  [  1,   2, 101, 101, 101]])
>>> inputs_tokens.shape
torch.Size([2, 5])

>>> cls_tokens = inputs_tokens == 101
tensor([[False,  True, False,  True, False],
        [False, False,  True,  True,  True]])

使用cls_tokens 掩码索引inputs_tokens 归结为减少 inputs_tokens 到cls_tokens 的true 值。在一般情况下，每批有不同数量的 true 值，保持形状是不可能的。

按照上面的例子，这里是seq_A：

>>> seq_A = torch.rand(2, 5, 1)
tensor([[[0.4644],
         [0.7656],
         [0.3951],
         [0.6384],
         [0.1090]],

        [[0.6754],
         [0.0144],
         [0.7154],
         [0.5805],
         [0.5274]]])

根据您的示例，您希望输出形状为(2, N, 1)。 N 会是什么？ 3?仅作为 2 true 值的第一批呢？生成的张量不能有不同的大小（axis=1 上的 2 和 3）。因此：“axis=1 上的所有序列必须具有相同的大小”。

---

但是，如果您希望每个批次具有相同数量的令牌 101，那么您可以通过广播您的索引张量而侥幸：

>>> inputs_tokens = torch.tensor([[  1, 101, 101, 101,   9],
                                  [  1,   2, 101, 101, 101]])
>>> inputs_tokens.shape

>>> N = cls_tokens[0].sum()
3

记住，我假设你有：

>>> assert all(cls_tokens.sum(axis=1) == N)

因此所需的输出（形状为(2, 3, 1)）是：

>>> seq_A[cls_tokens].reshape(seq_A.size(0), N, -1)
tensor([[[0.7656],
         [0.3951],
         [0.6384]],

        [[0.7154],
         [0.5805],
         [0.5274]]])

---

编辑 - 如果你真的想这样做，尽管你需要使用列表推导：

>>> [seq_A[i, cls_tokens[i]] for i in range(cls_tokens.size(0))]
[ tensor([[0.7656],
          [0.6384]]), 
  tensor([[0.7154],
          [0.5805],
          [0.5274]]) ]