collate_fn必须自定义以处理变长序列，因default_collate要求tensor形状一致；基础写法三步：pad、stack、封装为dict；padding_value需匹配tokenizer，batch_first=True确保(B,T)格式。
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_idslist 转成 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，按字段类型分组处理：
立即学习“Python免费学习笔记（深入）”；
- 序列类（需 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。
