Изменить форму и дополнить тензор с учетом длины

Question

Я дал двумерный тензор in формы a x b, подобный следующему (где a = 9 и каждый из A1, A2, ..., C2 представляет b -мерныйvector):

Кроме того, у меня есть массив lengths, где sum(lengths) = a, и каждая запись является положительным целым числом:

Затем я хотел бы получить трехмерный выходной тензор out, где первые lengths[0] записи in образуют первую строку, следующуюlengths[1] записи in образуют второй ряд и т. Д.То есть выходной тензор должен иметь форму len(lengths) x max(lengths) x b и дополняться нулями (каждый 0 на рисунке ниже представляет b -мерный нулевой вектор):

Поскольку это часть нейронной сети, которая обучается с использованием обратного распространения, все используемые операции должны быть дифференцируемыми .Как этого добиться (в идеале с хорошей производительностью ) с помощью PyTorch?

Answer 1

Вот моя реализация, использующая torch.nn.utils.rnn.pad_sequence():

in_tensor = torch.rand((9, 3))
print(in_tensor)
print(36*'=')
lengths = torch.tensor([3, 4, 2])
cum_len = 0
y = []
for idx, val in enumerate(lengths):
    y.append(in_tensor[cum_len : cum_len+val])
    cum_len += val
print(torch.nn.utils.rnn.pad_sequence(y, batch_first=True)))

вывод:

# 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]]])

Answer 2

Вы можете использовать функцию ниже. Он дифференцируем и может работать с 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]]])