如何在PyTorch中将句子长度批量转换为掩码？

Question

例如，来自

lens = [3, 5, 4]

我们想要得到

mask = [[1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1],
        [1, 1, 1, 1, 0]]

两者均为torch.LongTensor。

Answer 1

我发现的一种方法是：

torch.arange(max_len).expand(len(lens), max_len) < lens.unsqueeze(1)

如果有更好的方法，请分享！

Answer 2

只需对@ypc的答案提供一些解释（由于缺乏声誉而无法发表评论）：

torch.arange(max_len)[None, :] < lens[:, None]

总而言之，答案使用broadcasting机制隐式expand张量，就像在接受的答案中所做的那样。分步操作：

1. torch.arange（max_len）给您[0, 1, 2, 3, 4];

2. 添加[None, :]会在张量上附加第0个维度，使其形状为(1, 5)，从而为您提供[[0, 1, 2, 3, 4]];

3. 类似地，lens[:, None]将第一维附加到张量lens上，使其形状为(3, 1)，即[[3], [5], [4]]；

4. 通过遵循(1, 5)的规则，比较(3, 1)和broadcasting的张量（或进行类似+，-，*，/的操作）张量的形状为(3, 5)，结果值为result[i, j] = (j < lens[i])。

Answer 3

torch.arange(max_len)[None, :] < lens[:, None]