### 如何将 PyTorch 模型 (.pt 文件) 转换为 ONNX 格式 要将 PyTorch 的 `.pt` 或 `.pth` 文件转换为 ONNX 格式的模型，需遵循以下方法： #### 安装必要的依赖库 在开始之前，需要确保已安装 `onnx` 和 `onnxruntime` 库。可以通过以下命令完成安装[^3]： ```bash pip install onnx onnxruntime ``` #### 准备 PyTorch 模型 假设已经有一个训练好的 PyTorch 模型保存为 `.pt` 文件。首先需要加载该模型并将其设置为评估模式（即调用 `.eval()` 方法）。以下是加载模型的代码示例： ```python import torch # 加载预训练的 PyTorch 模型 model = torch.load("model.pt", map_location=torch.device('cpu')) model.eval() # 设置模型为评估模式 ``` #### 构造虚拟输入张量 为了导出模型至 ONNX 格式，需要提供一个虚拟输入张量作为样例输入。这个张量应具有与实际推理过程中一致的形状和数据类型。例如： ```python dummy_input = torch.randn(1, 3, 224, 224) # 假设输入尺寸为 (batch_size=1, channels=3, height=224, width=224) ``` #### 导出模型到 ONNX 格式 使用 `torch.onnx.export` 方法可实现从 PyTorch 到 ONNX 的转换。此函数接受多个参数，包括模型对象、虚拟输入张量以及目标 ONNX 文件路径等。具体代码如下所示[^1]: ```python # 导出模型到 ONNX 格式 torch.onnx.export( model, # 已加载的 PyTorch 模型 dummy_input, # 虚拟输入张量 "model.onnx", # 输出的 ONNX 文件名 export_params=True, # 是否导出训练后的权重参数 opset_version=11, # 所使用的 ONNX 版本号 do_constant_folding=True, # 是否执行常量折叠优化 input_names=['input'], # 输入节点名称列表 output_names=['output'], # 输出节点名称列表 dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} # 动态轴定义 ) ``` 上述代码中的关键参数解释如下： - **export_params**: 如果设置为 `True`，则会将训练后的权重参数嵌入到生成的 ONNX 文件中。 - **opset_version**: 表明所采用的 ONNX 操作集版本；推荐使用较新的稳定版（如 11 或更高）。 - **do_constant_folding**: 启动这一选项有助于减少最终 ONNX 图形大小并通过合并某些操作来提升性能。 - **dynamic_axes**: 当希望支持动态批量处理或其他维度变化时非常有用。 #### 验证导出的 ONNX 模型 一旦成功创建了 ONNX 文件，可以利用 `onnx.checker` 来验证其结构是否有效，并通过 `onnx.helper.printable_graph(onnx_model.graph)` 查看图形表示形式[^2]： ```python import onnx # 加载并检查 ONNX 模型的有效性 onnx_model = onnx.load("model.onnx") onnx.checker.check_model(onnx_model) print("ONNX Model Structure:") print(onnx.helper.printable_graph(onnx_model.graph)) ``` 最后，还可以借助 `onnxruntime` 对新生成的 ONNX 模型进行测试运行以确认功能正常。 --- ###