1. PyTorch与ONNX的相爱相杀为什么需要模型转换每次在PyTorch里跑通一个模型看着训练曲线蹭蹭上涨心里那个美啊但真要部署到生产环境时问题就来了——总不能要求客户也装个PyTorch吧这时候ONNX就像个救世主出现了。动态图 vs 静态图这事儿得好好唠唠。PyTorch的动态图就像用铅笔写草稿写错了随时能擦而ONNX的静态图则是把定稿的钢笔字拿去印刷改起来就费劲了。我去年给某家电企业做产线缺陷检测他们的嵌入式设备跑不动动态图就是靠转ONNX解决的。实测发现转换后的ResNet-18模型推理速度提升约23%CPU场景内存占用减少37%模型文件大小缩小15%注意模型转换不是银弹像带有控制流的LSTM模型转换时就容易踩坑。不过对于常见的CNN结构转换成功率能到90%以上。2. 手把手教你转换ResNet模型2.1 环境准备就像搭积木先来个全家桶安装建议用conda环境pip install torch2.0.1 onnx1.14.0 onnxruntime1.15.1 pip install netron # 可视化神器我习惯用ResNet做示例因为结构标准兼容性好社区资源丰富适合教学演示2.2 转换代码的魔鬼细节重点来看这个export函数import torch from torchvision.models import resnet18 model resnet18(pretrainedTrue).eval() # 必须设为eval模式 dummy_input torch.randn(1, 3, 224, 224) # 假输入要带batch维度 torch.onnx.export( model, dummy_input, resnet18.onnx, export_paramsTrue, # 必须True才能保存权重 opset_version13, # 版本太低会丢失算子 do_constant_foldingTrue, # 常量折叠优化 input_names[image], output_names[logits], dynamic_axes{ image: {0: batch}, # 动态batch维度 logits: {0: batch} } )踩坑记录去年用opset_version9转换MobileNetV3时silu激活函数直接消失dynamic_axes没设置对导致批量推理时崩溃忘记model.eval()batchnorm层全乱套3. 动态维度的魔法让模型灵活起来3.1 理解dynamic_axes参数这个字典决定了哪些维度可以变化dynamic_axes{ input: { 0: batch_size, # 第0维batch可变 2: height, # 图像高度可变 3: width # 图像宽度可变 } }实测案例对比设置方式支持输入形状推理耗时(ms)完全静态(1,3,224,224)45.2仅动态batch(N,3,224,224)46.8全动态输入尺寸(N,3,H,W)52.13.2 多输入/输出实战处理双输入模型比如图像文本# 假设模型定义 class MultiModalModel(torch.nn.Module): def forward(self, img, text): ... # 导出时要对应 torch.onnx.export( model, (dummy_img, dummy_text), # 元组形式传入 ..., input_names[image, text], output_names[output1, output2] )4. 转换后的模型体检验证与可视化4.1 模型健康检查三件套基础校验import onnx model onnx.load(resnet18.onnx) onnx.checker.check_model(model) # 抛出异常说明有问题推理验证import onnxruntime as ort ort_session ort.InferenceSession(resnet18.onnx) outputs ort_session.run( None, {image: np.random.randn(1,3,224,224).astype(np.float32)} )可视化神器Netronpython -m netron resnet18.onnx4.2 常见错误代码表错误类型解决方案Unsupported operator升级opset_version或自定义算子Shape inference failed检查dynamic_axes设置ValidationError用onnxruntime预加载测试精度下降超过1%检查模型是否处于eval模式5. 生产环境部署实战技巧5.1 ONNX Runtime性能调优在阿里云ECS上的测试数据ResNet50, batch16配置吞吐量(qps)延迟(ms)纯CPU78205CPUOpenMP125128CUDATensorRT42038启动优化会话options ort.SessionOptions() options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL options.execution_mode ort.ExecutionMode.ORT_SEQUENTIAL options.intra_op_num_threads 8 # 根据CPU核心数调整5.2 跨平台部署案例去年给某医疗客户做的部署方案训练端PyTorch 1.11 CUDA 11.3转换工具onnxruntime 1.10部署环境Windows医院工作站DirectMLLinux服务器CUDA安卓移动端NNAPI关键是要统一opset_version我们最终选用opset12作为公约数。6. 避坑指南来自血泪经验的建议版本兼容矩阵必须牢记PyTorch 1.8推荐用opset_version13ONNX Runtime最好和PyTorch同步更新自定义算子处理方案# 注册符号函数 torch.onnx.symbolic(custom_op) def symbolic_custom_op(g, input): return g.op(CustomOp, input) # 导出时添加注册表 torch.onnx.export(..., custom_opsets{custom_domain: 1})遇到诡异bug时的检查清单[ ] 模型是否调用了eval()[ ] dummy_input的dtype是否正确[ ] 输入输出名称是否匹配[ ] 是否开启了constant_folding最后说个真实案例有次转换总是失败折腾半天发现是模型里用了Python的random模块。记住——ONNX只认Tensor操作