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()，将序列后的mask中的所有条目设置为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]。
在每一行分别为lengths分配1，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可以给您完全想要的东西。

享受;）

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

1 个答案: