PyTorch `.pth` 转 ONNX：从模型训练到跨平台部署

<think>好的，用户想将预训练的PyTorch ResNet34模型（.pth文件）转换为ONNX格式。我需要根据提供的引用资料来整理步骤，并确保回答正确。 首先，用户提供的引用中有几个关于ResNet模型转换的例子，比如引用[2]和[4]提到了使用torch.onnx.export函数，以及设置export_params=True。引用[5]提到需要确定模型是保存了整个网络结构还是仅参数，这可能影响加载方式。 步骤应该是这样的：首先加载模型，确保模型结构和参数正确。然后准备输入张量，调用export函数。需要设置input_names和output_names，还有动态轴如果有需求。此外，引用[1]强调了export_params=True的重要性，必须包含进去。 需要注意用户可能遇到的常见问题，比如模型加载方式的不同（整个模型或仅参数），以及输入尺寸是否正确。还要验证ONNX模型是否正确，可以使用onnxruntime进行检查。 最后，生成相关问题，比如性能优化、部署应用等，参考之前的示例结构。</think>### 分步解答：PyTorch ResNet34模型转换为ONNX格式 #### 步骤1：加载预训练模型 1. **确定模型保存方式**： - 若`.pth`文件保存了完整模型结构+参数（如`torch.save(model, 'model.pth')`），直接加载： ```python model = torch.load('resnet34.pth') ``` - 若仅保存参数（如`torch.save(model.state_dict(), 'model.pth')`），需先实例化模型结构再加载权重： ```python model = resnet34() # 需提前导入ResNet34类定义 model.load_state_dict(torch.load('resnet34.pth')) ``` 引用[5]提到需测试模型加载方式以避免错误[^5]。 #### 步骤2：设置模型为推理模式 ```python model.eval() # 关闭Dropout/BatchNorm等训练专用层 ``` #### 步骤3：定义输入张量与转换函数 ```python import torch # 定义符合模型要求的输入尺寸（例如ResNet34常用224x224） dummy_input = torch.randn(1, 3, 224, 224) # (batch_size, channels, H, W) # 转换核心函数 torch.onnx.export( model, dummy_input, "resnet34.onnx", export_params=True, # 必须包含权重参数[^1] input_names=["input"], output_names=["output"], dynamic_axes={ "input": {0: "batch_size"}, # 支持动态batch_size "output": {0: "batch_size"} } ) ``` #### 步骤4：验证ONNX模型 使用`onnxruntime`检查模型是否有效： ```python import onnxruntime as ort ort_session = ort.InferenceSession("resnet34.onnx") outputs = ort_session.run(None, {"input": dummy_input.numpy()}) print("ONNX输出形状:", outputs[0].shape) # 应为(1, 1000) ``` #### 常见问题与优化 1. **输入尺寸匹配**：若训练时使用非标准输入（如96x96），需在转换时保持一致[^3]。 2. **动态维度**：通过`dynamic_axes`参数支持可变输入尺寸（如不同batch_size）。 3. **算子兼容性**：部分PyTorch算子需替换为ONNX支持的版本，可通过添加`opset_version=11`参数指定算子集版本。 ### 应用场景示例 - **跨平台部署**：ONNX模型可部署至TensorRT、OpenVINO等推理框架[^4]。 - **Web集成**：通过ONNX.js在浏览器中运行模型。