时间戳

首页 > 解决方案 > 在某个索引后用零填充火炬张量

python - 在某个索引后用零填充火炬张量

问题描述

给定一个 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

标签： pythonnlppytorch

解决方案

您可以使用二进制掩码来执行此操作。
使用lengths列索引来mask指示每个序列的结束位置（请注意，我们使mask长度超过a.size(1)允许全长序列的长度）。
使用我们将seq len 之后的cumsum()所有条目设置为 1。mask

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 ，如下所示：lengthsmask

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会给你你想要的。

享受 ;）

推荐阅读