python - 在特定索引后用零填充 torch 张量

Question

给定一个 3d 张量，说: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() 我们将所有条目设置为 mask在 seq 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给你你想要的。
享受 ;)