本文介绍如何高效合并多个结构相同（键一致）、值为 PyTorch 张量的字典，替代手动遍历+defaultdict+torch.stack的冗余流程，推荐直接复用 torch.utils.data.default_collate——它专为此类结构化批量聚合而设计，简洁、健壮且性能更优。
本文介绍如何高效合并多个结构相同（键一致）、值为 pytorch 张量的字典，替代手动遍历+`defaultdict`+`torch.stack`的冗余流程，推荐直接复用 `torch.utils.data.default_collate`——它专为此类结构化批量聚合而设计，简洁、健壮且性能更优。
在 PyTorch 数据处理中，常需将一批样本（每个样本为一个键值映射字典，如 {'input_ids': tensor(...), 'attention_mask': tensor(...)}）合并为统一的批数据字典。传统做法是手动收集各键对应的张量列表，再逐个调用 torch.stack，代码冗长且易出错：
from collections import defaultdict
import torch
mention_inputs = defaultdict(list)
for idx in mention_indices:
mention_input, _ = get_mention_sample(idx) # 假设返回 dict[str, Tensor]
for key, value in mention_input.items():
mention_inputs[key].append(value)
# 合并为 batched dict: {key → (B, ...) tensor}
mention_inputs = {k: torch.stack(v) for k, v in mention_inputs.items()}该方法虽可行，但存在明显缺陷：
- 需显式初始化 defaultdict 并双重循环；
- 若某键值非张量（如标量或嵌套结构），torch.stack 会报错，缺乏容错性；
- 未利用 PyTorch 内置的成熟批处理逻辑，重复造轮子。
✅ 更优解：直接使用 torch.utils.data.default_collate
该函数是 PyTorch DataLoader 的默认批合并器，专为“同结构字典/列表/元组的张量聚合”优化，支持自动递归处理嵌套结构、类型推断与形状对齐：
from torch.utils.data.dataloader import default_collate
# mention_indices 是 list[dict]，每个 dict 键相同、值为 shape-(D) 的 tensor
batched_dict = default_collate(mention_indices)
# 输出示例: {'input_ids': tensor(B, D), 'attention_mask': tensor(B, D)}default_collate 的核心优势在于：
? 自动结构识别：检测输入是否为字典列表，并按键分组；
? 智能堆叠：对同键下所有张量执行 torch.stack(..., dim=0)；
? 扩展兼容：天然支持嵌套字典、列表、元组及混合类型（如 {'x': tensor, 'y': int} → 'y' 被转为 tensor([int1, int2, ...])）；
? 生产就绪：经大量训练场景验证，异常处理完善（如维度不一致时抛出清晰错误）。
⚠️ 注意事项：
- 所有字典必须键集合完全一致，缺失键会导致 KeyError；
- 同一键下的所有张量除 batch 维外形状必须严格相同（如 (D,)、(H, W)），否则 stack 失败；
- 若需自定义行为（如 padding 变长序列），应继承 default_collate 或实现专用 collate_fn，而非绕过它。
总结：在构建自定义数据流（如 mention-level 特征聚合）时，优先选用 default_collate 替代手写合并逻辑——它更简洁、更鲁棒、更符合 PyTorch 生态惯例，且零额外依赖。
