参考文章:https://pytorch.org/docs/stable/generated/torch.nn.utils.rnn.pad_sequence.html
首先看参考文章里的解释,如下图
从文章描述来看,当 batch_first = True 时,输出的张量的 size 是 B x T x *。当 batch_first = False 时,输出的张量的 size 是 T x B x *
一个例子如下代码:
seqs = [torch.tensor([1,2,3]), torch.tensor([3,4])] # Sequences
seq_lens = torch.tensor([3,2]) # Actual lengths of sequences
# First, pad the sequences to the same length
# 在处理序列数据时,由于序列的长度可能不一致,为了能够将它们组织成一个张量以便进行批处理,
# 通常需要对这些序列进行填充,使它们的长度保持一致。pad_sequence 函数的作用就是对一批序列进行填充,
# 使得它们的长度都与最长序列的长度相同。
padded_seqs = nn.utils.rnn.pad_sequence(seqs, batch_first=True)
print(padded_seqs)
print(padded_seqs.size())
padded_seqs = nn.utils.rnn.pad_sequence(seqs, batch_first=False)
print(padded_seqs)
print(padded_seqs.size())
输出日志如下:
tensor([[1, 2, 3],
[3, 4, 0]])
torch.Size([2, 3])
tensor([[1, 3],
[2, 4],
[3, 0]])
torch.Size([3, 2])
正好映证了 文章里的解释