PyTorch中如何将不同长度序列打包为批次并忽略填充部分

1. 序列填充实现 #

可变长度序列的填充推荐在DataLoader的collate_fn中实现，可按每个批次的最大序列长度动态填充，减少冗余计算，解决默认拼接逻辑要求样本维度一致的报错。
具体实现代码：

from torch.utils.data import DataLoader
from torch.nn.utils.rnn import pad_sequence
import torch
import torch.nn as nn
import torch.nn.functional as F

def collate_fn(batch):
    # batch为Dataset.__getitem__返回的样本列表，每个样本结构为(seq_tensor, label)，其中seq_tensor维度为(x, 128)
    seqs = [item[0] for item in batch]
    labels = [item[1] for item in batch]
    # 记录每个序列的原始长度
    seq_lens = torch.tensor([len(s) for s in seqs], dtype=torch.long)
    # 执行填充，默认补0，可通过padding_value参数自定义填充值
    padded_seqs = pad_sequence(seqs, batch_first=True) # 输出维度为(batch_size, max_len, 128)
    return padded_seqs, seq_lens, torch.tensor(labels)

# DataLoader传入自定义collate_fn即可完成批次自动填充
dataloader = DataLoader(your_dataset, batch_size=32, collate_fn=collate_fn)

2. 填充内容过滤实现 #

通过有效位置掩码（mask）即可实现填充内容的过滤，全程采用张量运算，不需要编写Python循环，效率和原生算子一致。

2.1 前馈子网络批量计算 #

对每个序列元素独立计算的nn.Sequential子网络，可直接对填充后的张量做维度转换后批量计算：

class SeqModel(nn.Module):
    def __init__(self, feature_dim=128, hidden_dim=64, out_dim=1):
        super().__init__()
        self.feature_extractor = nn.Sequential(
            nn.Linear(feature_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, out_dim)
        )
    
    def forward(self, padded_seqs, seq_lens):
        batch_size, max_len, feat_dim = padded_seqs.shape
        # 维度转换为(batch_size*max_len, feat_dim)，批量过子网络
        flat_feats = padded_seqs.reshape(-1, feat_dim)
        flat_outputs = self.feature_extractor(flat_feats)
        # 恢复维度为(batch_size, max_len, out_dim)
        seq_outputs = flat_outputs.reshape(batch_size, max_len, -1)
        
        # 生成有效位置掩码，维度为(batch_size, max_len)，有效位置为True，填充位置为False
        mask = torch.arange(max_len, device=seq_lens.device)[None, :] < seq_lens[:, None]
        mask = mask.unsqueeze(-1) # 扩展维度和输出对齐，方便后续计算
        
        # 填充位置输出置0，不影响后续累加/平均计算
        seq_outputs = seq_outputs * mask
        
        # 示例：生成样本级输出，可选求和/平均
        # 求和实现：填充位置为0，无额外影响
        sample_sum_output = seq_outputs.sum(dim=1)
        # 平均实现：除以对应样本的有效长度
        # sample_avg_output = seq_outputs.sum(dim=1) / seq_lens.unsqueeze(-1)
        
        return sample_sum_output, seq_outputs, mask

2.2 损失计算忽略填充部分 #

如果是序列元素级损失，通过掩码过滤填充位置的损失值即可：

def sequence_loss(seq_outputs, target, mask):
    # 计算无归约的损失
    loss = F.mse_loss(seq_outputs, target, reduction='none')
    # 过滤填充位置损失
    loss = loss * mask
    # 按有效元素数量做归一化
    avg_loss = loss.sum() / mask.sum()
    return avg_loss

方案说明 #

该方案完全适配前馈网络场景，不需要依赖RNN类模块的内置填充处理能力，所有运算均支持自动微分，梯度计算和反向传播不受填充内容影响。

内容的提问来源于stack exchange，提问作者Michael