在特定索引后用零填充火炬张量

Question

给定一个 3d Tenzor，说： batch x sentence length x embedding dim

a = torch.rand((10, 1000, 96)) 

以及每个句子的实际长度数组（或张量）

lengths =  torch .randint(1000,(10,))

outputs tensor([ 370., 502., 652., 859., 545., 964., 566., 576.,1000., 803.])

如何根据张量“长度”在维度 1（句子长度）的某个索引后用零填充张量“a”？

我想要这样的：

a[ : , lengths : , : ]  = 0

一种方法（如果批量足够大，则速度较慢）：

for i_batch in range(10):
    a[ i_batch  , lengths[i_batch ] : , : ]  = 0

Answer 1

您可以使用二进制掩码来完成。
使用lengths的列索引来mask我们指明每个序列结束（注意，我们做mask超过a.size(1)允许与全长序列）。
使用cumsum()我们将maskseq len 之后的所有条目设置为 1。

mask = torch.zeros(a.shape[0], a.shape[1] + 1, dtype=a.dtype, device=a.device)
mask[(torch.arange(a.shape[0]), lengths)] = 1
mask = mask.cumsum(dim=1)[:, :-1]  # remove the superfluous column
a = a * (1. - mask[..., None])     # use mask to zero after each column

对于a.shape = (10, 5, 96)， 和lengths = [1, 2, 1, 1, 3, 0, 4, 4, 1, 3]。
将 1 分配给lengths每行的相应值，mask如下所示：

mask = 
tensor([[0., 1., 0., 0., 0., 0.],
        [0., 0., 1., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0.],
        [0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0.],
        [1., 0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 1., 0.],
        [0., 0., 0., 0., 1., 0.],
        [0., 1., 0., 0., 0., 0.],
        [0., 0., 0., 1., 0., 0.]])

之后cumsum你

mask = 
tensor([[0., 1., 1., 1., 1.],
        [0., 0., 1., 1., 1.],
        [0., 1., 1., 1., 1.],
        [0., 1., 1., 1., 1.],
        [0., 0., 0., 1., 1.],
        [1., 1., 1., 1., 1.],
        [0., 0., 0., 0., 1.],
        [0., 0., 0., 0., 1.],
        [0., 1., 1., 1., 1.],
        [0., 0., 0., 1., 1.]])

请注意，它在有效序列条目所在的位置准确地具有零，并且在序列长度之外具有一个。采取1 - mask给你正是你想要的。

享受 ;）