本文介绍如何高效地将多个结构相同（键一致、值为同形张量）的字典合并为一个批量字典，避免手动循环+torch.stack，推荐直接复用 PyTorch 内置的 default_collate 函数，兼顾简洁性、鲁棒性与性能。
本文介绍如何高效地将多个结构相同（键一致、值为同形张量）的字典合并为一个批量字典，避免手动循环+`torch.stack`，推荐直接复用 pytorch 内置的 `default_collate` 函数，兼顾简洁性、鲁棒性与性能。
在 PyTorch 数据处理流程中，常需将一批样本（每个样本为 dict[str, Tensor]）按字段聚合为批量字典——例如，{'input_ids': [t1, t2, ..., tn], 'attention_mask': [m1, m2, ..., mn]} → {'input_ids': torch.stack([t1,t2,...,tn]), 'attention_mask': torch.stack([m1,m2,...,mn])}。虽然手动遍历 + defaultdict + torch.stack 可行，但存在冗余逻辑、类型校验缺失及扩展性不足等问题。
更优解是直接使用 PyTorch 官方提供的 default_collate：
from torch.utils.data.dataloader import default_collate
# 假设 mention_inputs 是一个包含 N 个字典的列表，每个字典结构相同：
# [
# {'feat_a': tensor([1.0, 2.0]), 'feat_b': tensor([0.1])},
# {'feat_a': tensor([3.0, 4.0]), 'feat_b': tensor([0.2])},
# ...
# ]
batch_dict = default_collate(mention_inputs)
# 输出：{'feat_a': tensor([[1., 2.], [3., 4.]]), 'feat_b': tensor([[0.1], [0.2]])}default_collate 是 DataLoader 默认批处理函数，专为结构化数据设计：它递归识别嵌套结构，对同键张量自动执行 torch.stack（要求所有张量 shape 完全一致），对非张量（如 int/float/str）则转为 torch.tensor；若遇到不兼容类型（如混合张量与 list），会抛出清晰错误，便于调试。
✅ 优势总结：
- 零代码封装：无需手写 defaultdict 或显式循环；
- 健壮性强：内置 shape 校验、设备/梯度一致性检查（如混合 CPU/Tensor 会报错）；
- 开箱即用：天然支持嵌套字典、列表、元组等常见结构；
- 性能可靠：底层经充分优化，与 DataLoader 同源，无额外开销。
⚠️ 注意事项：
- 所有同键张量必须具有完全相同的 shape 和 dtype，否则 default_collate 将报 RuntimeError；
- 若需自定义行为（如 padding 变长序列），应继承并重写 collate_fn，而非绕过 default_collate；
- 不适用于含不可序列化对象（如 lambda、文件句柄）的字典——这本就不符合 PyTorch 数据流规范。
综上，当面对“同键字典列表 → 批量张量字典”这一高频任务时，default_collate 不仅是最简方案，更是最符合 PyTorch 设计哲学的工程实践。
