PyTorch DataLoader迭代报TypeError及大样本模型测试咨询

根因定位 #

报错直接触发原因是类型不匹配：

• 代码中MAX_LEN='512'赋值为字符串类型，而tokenizer在做长度截断校验时，会将计算得到的文本token长度（整数类型）和传入的max_length参数做大小比较，整数和字符串无法直接比较，触发TypeError。
• 代码中还存在两个潜在隐患：一是标签target传入的是字符串格式的'0'/'1'，后续转张量时存在隐式转换风险；二是使用的pad_to_max_length参数在新版Transformers库中已被废弃，会触发兼容警告。

修复方案 #

1. 修正MAX_LEN的类型，去掉单引号改为整数赋值：

MAX_LEN = 512

2. 提前将标签列转为整数类型，避免隐式转换问题：

df['target'] = df['target'].astype(int)

3. 替换废弃的分词参数，兼容新版Transformers库，将__getitem__方法中的分词调用修改为：

encoding = self.tokenizer.encode_plus(
    Paragraph,
    add_special_tokens=True,
    max_length=self.max_len,
    return_token_type_ids=False,
    padding='max_length', # 替换原pad_to_max_length=True
    return_attention_mask=True,
    return_tensors='pt',
)

针对3M规模的测试场景，从数据加载、推理执行全链路优化，可实现数倍到数十倍的效率提升：

• 前置分词处理，消除实时计算开销
  不要在__getitem__方法中做实时分词，这是DataLoader加载慢的最核心原因。提前批量对全量300万样本完成分词，将处理好的input_ids、attention_mask、targets存为numpy数组、内存映射（memmap）文件或者Apache Arrow格式，后续Dataset直接读取预处理完成的张量，完全跳过实时分词的CPU开销，加载速度可提升5~10倍。
• DataLoader参数调优
  • num_workers根据CPU物理核心数调整，通常设置为核心数的70%~100%即可，不要硬编码为4，避免加载速度跟不上GPU计算速度；
  • 开启pin_memory=True，使用锁页内存加速CPU到GPU的张量传输；
  • PyTorch1.7及以上版本开启persistent_workers=True，避免每个epoch重复初始化worker进程，减少额外开销；
  • 在显存允许范围内尽可能调大batch_size，不要设置为1这类极小值，避免GPU利用率过低。
    优化后的DataLoader初始化参考：

  return DataLoader(
      ds,
      batch_size=batch_size,
      num_workers=8, # 按自身CPU核心数调整
      pin_memory=True,
      persistent_workers=True
  )
  

• 存储格式优化
  不要用CSV、普通Pandas DataFrame存储300万量级样本，这类格式顺序读取速度慢、内存占用高。替换为Apache Arrow、HDF5或者numpy memmap格式，配合HuggingFace Datasets库可实现零拷贝内存映射读取，不需要把全量数据加载到内存，普通消费级机器也能顺畅处理3M样本，读取速度比CSV快一个数量级。
• 推理阶段优化
  • 模型切换到评估模式，使用torch.inference_mode()（比torch.no_grad()速度更快）包裹推理逻辑，关闭梯度计算，减少显存占用和计算开销；
  • 支持混合精度的GPU开启torch.cuda.amp.autocast()做混合精度推理，速度可提升1倍左右，精度损失可忽略；
  • 固定batch大小的批量测试场景，可以用ONNX Runtime、TensorRT对模型做编译优化，推理吞吐量可再提升2~4倍。
• 裁剪无效负载
  测试阶段不要在DataLoader返回结果中携带原始文本等不需要输入模型的字段，减少进程间数据序列化、传输的开销，需要输出结果关联原始文本时，可通过索引单独关联。

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