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

Question

例如，来自

lens = [3, 5, 4]

我们想得到

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

两者都是torch.LongTensors。

Answer 1

只是对@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 2

我发现的一种方法是：

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

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