如何 return 基于 1D mask 截断的 k-dim pytorch 张量

Question

所以假设我有一个 k-dim 张量和一个 1-dim 掩码，它在 pytorch 中经常用于可变长度序列，我想 return 一个代表最多元素的张量掩码中的第一个假值。这是一个例子：

import torch

a = torch.tensor([[1,2],[3,4],[5,6],[0,0],[0,0],[0,0]])
b = torch.tensor([True,True,True,False,False,False])

# magic goes here, result of c should be:

print(c)
>>> [[1,2],[3,4],[5,6]]

在此示例中，输入张量是二维的，但它可以是 k-d，在这些维度上具有任意数量的值。只有第一个维度需要匹配掩码维度。所以做 torch.masked_select 是行不通的，因为要截断的张量不像掩码那样是一维的，并且由于您不知道维数，所以压缩和取消压缩也不是解决方案。

掩码对于前 k 个元素始终为真，对于其余元素始终为假，但如果您的解决方案不“依赖”于此，那也没关系。

这似乎是人们一直在做的事情，但我找不到任何地方回答过这个问题。

Answer 1

您可以简单地将掩码作为切片索引传递给张量：

c = a[b]

>>> c
tensor([[1, 2],
        [3, 4],
        [5, 6]])

如何 return 基于 1D mask 截断的 k-dim pytorch 张量

How to return a k-dim pytorch tensor truncated based on 1D mask

python

boolean

pytorch