Pytorch 强制 pad_sequence 到特定长度

在本文中，我们将介绍如何使用PyTorch中的pad_sequence函数将序列填充到特定的长度。pad_sequence函数是PyTorch中的一个实用工具，可用于将批量序列中的元素通过填充操作转换为具有相同长度的序列。这在处理自然语言处理（NLP）任务中特别有用，例如文本分类、机器翻译等。

阅读更多：Pytorch 教程

pad_sequence函数介绍

pad_sequence函数位于torch.nn.utils模块中，用于将一批序列进行填充操作，使它们具有相同的长度。该函数的定义如下：

torch.nn.utils.rnn.pad_sequence(sequences, batch_first=False, padding_value=0.0)

其中：
– sequences是一个列表，包含了需要填充的序列。
– batch_first参数决定了填充后的张量的形状。如果为True，则返回的张量形状为(batch_size, max_sequence_length, sequence_dim)；如果为False（默认值），则返回的张量形状为(max_sequence_length, batch_size, sequence_dim)。
– padding_value参数指定了填充的值，默认为0.0。

示例说明

让我们以一个简单的例子来说明如何使用pad_sequence函数将序列填充到特定的长度。

假设我们有一个包含5个序列的列表，每个序列的长度不一样：

import torch
from torch.nn.utils.rnn import pad_sequence

sequences = [torch.tensor([1, 2, 3]),
             torch.tensor([4, 5]),
             torch.tensor([6, 7, 8, 9]),
             torch.tensor([10]),
             torch.tensor([11, 12, 13, 14, 15])]

现在，我们希望将这些序列填充到最长序列的长度，即5。我们可以通过调用pad_sequence函数来实现：

padded_sequences = pad_sequence(sequences, padding_value=0)

执行完上述代码后，我们得到了填充后的张量padded_sequences：

tensor([[ 1,  2,  3,  0,  0],
        [ 4,  5,  0,  0,  0],
        [ 6,  7,  8,  9,  0],
        [10,  0,  0,  0,  0],
        [11, 12, 13, 14, 15]])

可以看到，所有的序列都被填充到了长度为5。未填充的位置被用0填充。

如果我们将batch_first参数设置为True，函数将返回以批次为主的张量形状：

padded_sequences_batch_first = pad_sequence(sequences, batch_first=True, padding_value=0)

执行上述代码后，我们得到了批次为主的张量padded_sequences_batch_first：

tensor([[ 1,  2,  3,  0],
        [ 4,  5,  0,  0],
        [ 6,  7,  8,  9],
        [10,  0,  0,  0],
        [11, 12, 13, 14],
        [ 0,  0,  0, 15]])