PyTorch自定义collate_fn处理变长序列Batch方法

本文深入解析了PyTorch中处理变长序列数据的关键环节——自定义collate_fn的必要性与实战技巧：从为何default_collate在文本、语音等场景下必然失效，到三步核心写法（pad→stack→封装）、padding_value与batch_first的精准配置陷阱，再到多字段协同处理（input_ids/attention_mask/ner_tags等）的工程化组织策略，强调collate_fn应专注“无脑填充”，而截断、对齐等逻辑必须前置到dataset中，帮助读者避开常见维度错、loss不降、attention分心等致命坑，真正写出健壮、可维护、与模型严丝合缝的数据组装逻辑。

collate_fn 为什么必须自己写？

PyTorch 的 DataLoader 默认用 default_collate，它要求 batch 内所有样本的 tensor 形状完全一致。一旦遇到变长序列（比如不同长度的文本、语音帧、轨迹点），default_collate 直接报错：RuntimeError: stack expects each tensor to be equal size。这不是数据有问题，是拼接逻辑没适配——你得接管 batch 组装过程。

怎么写一个基础版 collate_fn 处理变长文本？

核心就三步：pad、stack、封装成 dict/tuple。以 NLP 场景为例，假设每个样本是 {'input_ids': [12, 45, 67], 'label': 1} 这种不定长 list，你需要：

• 用 torch.nn.utils.rnn.pad_sequence 对所有 input_ids list 转成 tensor 后做右填充（注意传入的是 list[torch.Tensor]，不是 list[list]）
• 把 label 直接堆叠成 torch.stack（它本身定长）
• 返回字典，保持和 dataset __getitem__ 返回结构一致

def collate_fn(batch):
    input_ids = [torch.tensor(item['input_ids']) for item in batch]
    labels = torch.tensor([item['label'] for item in batch])
    padded = torch.nn.utils.rnn.pad_sequence(
        input_ids, batch_first=True, padding_value=0
    )
    return {'input_ids': padded, 'label': labels}

padding_value 和 batch_first 容易踩哪些坑？

padding_value 必须和你的模型 embedding 层对齐：BERT 类模型通常用 0（对应 [PAD] token id），但有些自定义 tokenizer 可能用 -1 或 100；batch_first=True 是为了输出 shape 为 (B, T)，否则是 (T, B)，后续送进 Transformer 会出维度错。

• 如果用了 batch_first=False 却没调换模型输入维度，会触发 Expected input to have 3 dimensions, got 2
• padding_value 错了会导致模型把 pad 当作真实 token 学习，loss 不降、attention 分心
• 别在 collate_fn 里做截断（如 [:max_len]）——该操作应放在 dataset 的 __getitem__ 中，避免 collate_fn 逻辑膨胀

多字段 + 不同 padding 策略怎么组织？

真实任务常有 input_ids、attention_mask、token_type_ids 甚至 ner_tags（也变长）。不要硬写 if-else，按字段类型分组处理：

• 序列类（需 pad）：统一用 pad_sequence(..., padding_value=X)，各自指定 padding_value
• 标量类（label、uid）：用 torch.tensor(...) 直接转
• mask 类（attention_mask）：和 input_ids 同长度，pad 时用 0 补，保证 pad 位置 mask=0
• 别忘了返回的 key 名必须和模型 forward 参数名严格一致，否则 model(**batch) 会报 unexpected keyword argument

复杂度真正上来的地方不在 padding，而在对齐：比如 ner_tags 长度要和 input_ids 对齐，但又不能简单 pad 到同一 max_len——因为 subword 分词后，tag 序列可能比原始 word 更长。这时候得在 dataset 里完成对齐，collate_fn 只负责无脑 pad。

文中关于的知识介绍，希望对你的学习有所帮助！若是受益匪浅，那就动动鼠标收藏这篇《PyTorch自定义collate_fn处理变长序列Batch方法》文章吧，也可关注golang学习网公众号了解相关技术文章。