将PyTorch模型转ONNX生成多文件的原因排查求助

可能的原因 #

• 大张量触发外部数据存储：当模型中存在尺寸超过默认阈值（2GB）的张量时，PyTorch的ONNX导出逻辑会自动将大权重拆分为外部辅助文件存储，主.onnx文件仅保留引用路径。你的模型包含BERT预训练权重，大概率触发了这个机制。
• 自定义模型结构的导出兼容性问题：如果EmotionModel作为多模态模型，内部子模块的权重未被完全整合到主计算图中，导出时会将分离的参数单独存储为额外文件。
• 旧版本PyTorch的导出参数默认行为：部分旧版PyTorch中，export_params或do_constant_folding的默认设置可能导致参数拆分存储。

解决办法 #

1. 显式禁用外部数据存储
   在torch.onnx.export中添加参数，强制将所有参数嵌入主ONNX文件（需PyTorch 1.10及以上版本支持save_as_external_data参数）：

   torch.onnx.export(model,
                     (input_ids_tf,token_type_ids_tf,attention_mask_tf,feature_wav_tf),
                     onnx_name,
                     opset_version=11,
                     input_names=['input_ids', 'token_type_ids', 'attention_mask', 'feature_wav'],
                     output_names=['logits'],
                     dynamic_axes={"input_ids": {0: "batch_size"},
                                   "token_type_ids": {0: "batch_size"},
                                   "attention_mask": {0: "batch_size"},
                                   "feature_wav": {0: "batch_size"},},
                     export_params=True,
                     do_constant_folding=True,
                     save_as_external_data=False  # 强制所有参数嵌入主文件
                     )
   

2. 验证模型参数完整性
   导出前检查所有参数是否属于主模型实例，避免子模块权重分离：

   # 打印所有参数的所属模块与尺寸
   for name, param in model.named_parameters():
       print(f"{name}: {param.shape}")
   

   若存在未整合的子模块，需调整模型加载逻辑，确保所有权重被正确加载到主模型中。

3. 调整ONNX opset版本或PyTorch版本
   尝试切换到opset 13/14版本，或升级PyTorch到最新稳定版，修复旧版本导出的兼容性bug。

内容的提问来源于stack exchange，提问作者Frank.Fan