Pytorch 根据给定的长度列表重塑并填充张量

我在中给出了一个二维张量

，其形状为
a x b
，如下所示（其中
a=9
和
A1
，
A2
，…，
C2
表示一个
b
-维向量）：


此外，我有一个长度为
的数组，其中
总和（长度）=a
，每个条目都是一个正整数：


然后我想获得一个3d输出张量
out
，其中
中
的第一个
长度[0]
条目构成第一行，下一个
长度[1]
条目构成第二行，依此类推。也就是说，输出张量应具有形状
len（长度）x max（长度）x b
，并用零填充（下图中的每个
0
表示
b
维零向量）：


由于这是使用反向传播训练的神经网络的一部分，因此使用的所有操作都必须是可微的。如何使用Pytork实现这一点（理想情况下，性能良好？
您可以使用以下功能。它是可微的，可以与backprop一起使用

def sequence_to_padding(x, length): 
    # declare the shape, it can work for x of any shape.
    ret_tensor = torch.zeros((length.shape[0], torch.max(length)) + tuple(x.shape[1:])) 
    cum_len = 0  
    for i, l in enumerate(length): 
        ret_tensor[i, :l] = x[cum_len: cum_len+l] 
        cum_len += l 
    return ret_tensor 

例如：

in_vector = torch.rand((9,1))  
#tensor([[0.3545],
#    [0.5443],
#    [0.7550],
#    [0.9624],
#    [0.9250],
#    [0.8035],
#    [0.6877],
#    [0.4186],
#    [0.4199]])
lengths = torch.tensor([3, 4, 2])  
sequence_to_padding(in_vector, lengths)
#tensor([[[0.3545],
#     [0.5443],
#     [0.7550],
#     [0.0000]],
#
#    [[0.9624],
#     [0.9250],
#     [0.8035],
#     [0.6877]],
#
#    [[0.4186],
#     [0.4199],
#     [0.0000],
#     [0.0000]]])

您可以使用下面的功能。它是可微的，可以与backprop一起使用

def sequence_to_padding(x, length): 
    # declare the shape, it can work for x of any shape.
    ret_tensor = torch.zeros((length.shape[0], torch.max(length)) + tuple(x.shape[1:])) 
    cum_len = 0  
    for i, l in enumerate(length): 
        ret_tensor[i, :l] = x[cum_len: cum_len+l] 
        cum_len += l 
    return ret_tensor 

例如：

in_vector = torch.rand((9,1))  
#tensor([[0.3545],
#    [0.5443],
#    [0.7550],
#    [0.9624],
#    [0.9250],
#    [0.8035],
#    [0.6877],
#    [0.4186],
#    [0.4199]])
lengths = torch.tensor([3, 4, 2])  
sequence_to_padding(in_vector, lengths)
#tensor([[[0.3545],
#     [0.5443],
#     [0.7550],
#     [0.0000]],
#
#    [[0.9624],
#     [0.9250],
#     [0.8035],
#     [0.6877]],
#
#    [[0.4186],
#     [0.4199],
#     [0.0000],
#     [0.0000]]])

下面是我使用
torch.nn.utils.rnn.pad\u sequence（）实现的：

输出：

# in_tensor of shape (9 x 3)
tensor([[0.9169, 0.3549, 0.6211],
        [0.4832, 0.5475, 0.8862],
        [0.8708, 0.5462, 0.9374],
        [0.4605, 0.1167, 0.5842],
        [0.1670, 0.2862, 0.0378],
        [0.2438, 0.5742, 0.4907],
        [0.1045, 0.5294, 0.5262],
        [0.0805, 0.2065, 0.2080],
        [0.6417, 0.4479, 0.0688]])
====================================
# out tensor of shape (len(lengths) x max(lengths) x b), in this case b is 3
tensor([[[0.9169, 0.3549, 0.6211],
         [0.4832, 0.5475, 0.8862],
         [0.8708, 0.5462, 0.9374],
         [0.0000, 0.0000, 0.0000]],

        [[0.4605, 0.1167, 0.5842],
         [0.1670, 0.2862, 0.0378],
         [0.2438, 0.5742, 0.4907],
         [0.1045, 0.5294, 0.5262]],

        [[0.0805, 0.2065, 0.2080],
         [0.6417, 0.4479, 0.0688],
         [0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000]]])

下面是我使用
torch.nn.utils.rnn.pad\u sequence（）实现的：

输出：

# in_tensor of shape (9 x 3)
tensor([[0.9169, 0.3549, 0.6211],
        [0.4832, 0.5475, 0.8862],
        [0.8708, 0.5462, 0.9374],
        [0.4605, 0.1167, 0.5842],
        [0.1670, 0.2862, 0.0378],
        [0.2438, 0.5742, 0.4907],
        [0.1045, 0.5294, 0.5262],
        [0.0805, 0.2065, 0.2080],
        [0.6417, 0.4479, 0.0688]])
====================================
# out tensor of shape (len(lengths) x max(lengths) x b), in this case b is 3
tensor([[[0.9169, 0.3549, 0.6211],
         [0.4832, 0.5475, 0.8862],
         [0.8708, 0.5462, 0.9374],
         [0.0000, 0.0000, 0.0000]],

        [[0.4605, 0.1167, 0.5842],
         [0.1670, 0.2862, 0.0378],
         [0.2438, 0.5742, 0.4907],
         [0.1045, 0.5294, 0.5262]],

        [[0.0805, 0.2065, 0.2080],
         [0.6417, 0.4479, 0.0688],
         [0.0000, 0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000]]])