Pytorch的packed_sequence / pad_sequence垂直填充张量，用于张量列表

Question

I am trying to pad sequence of tensors for LSTM mini-batching, where each timestep in the sequence contains a sub-list of tensors (representing multiple features in a single timestep).

For example, sequence 1 would have 3 timesteps and within each timestep there are 2 features. An example below would be:

Sequence 1 = [[1,2],[2,2],[3,3],[3,2],[3,2]]

Sequence 2 = [[4,2],[5,1],[4,4]]

Sequence 3 = [[6,9]]

I run pytorch's pad_sequence function (this goes for pack_sequence too) like below:

import torch
import torch.nn.utils.rnn as rnn_utils

a = torch.tensor([[1,2],[2,2],[3,3],[3,2],[3,2]])
b = torch.tensor([[4,2],[5,1],[4,4]])
c = torch.tensor([[6,9]])
result = rnn_utils.pad_sequence([a, b, c])

My expected output is as follows:

Sequence 1 = [[1,2],[2,2],[3,3],[3,2],[3,2]]

Sequence 2 = [[4,2],[5,1],[4,4],[0,0],[0,0]]

Sequence 3 = [[6,9],[0,0],[0,0],[0,0],[0,0]]

However, the output I got is as follows:

tensor([[[1, 2],
         [4, 2],
         [6, 9]],

        [[2, 2],
         [5, 1],
         [0, 0]],

        [[3, 3],
         [4, 4],
         [0, 0]],

        [[3, 2],
         [0, 0],
         [0, 0]],

        [[3, 2],
         [0, 0],
         [0, 0]]])

The padding seems to go vertically rather than what I expect. How do I go about getting the correct padding that I need?

Answer 1

Simply change

result = rnn_utils.pad_sequence([a, b, c])

to

result = rnn_utils.pad_sequence([a, b, c], batch_first=True)
seq1 = result[0]
seq2 = result[1]
seq3 = result[2]

By default, batch_first is False. Output will be in B x T x * if True, or in T x B x * otherwise, where

B is batch size. It is equal to the number of elements in sequences ,

T is length of the longest sequence, and

* is any number of trailing dimensions, including none.

output:

tensor([[1, 2],
        [2, 2],
        [3, 3],
        [3, 2],
        [3, 2]]) # sequence 1
tensor([[4, 2],
        [5, 1],
        [4, 4],
        [0, 0],
        [0, 0]]) # sequence 2
tensor([[6, 9],
        [0, 0],
        [0, 0],
        [0, 0],
        [0, 0]]) # sequence 3