### 将 PyTorch `.pt` 模型文件以 `eval` 模式转换为 ONNX 格式的实现思路及完整 Python 脚本 将 PyTorch 的 `.pt` 模型文件导出为 ONNX 格式的过程涉及以下几个核心步骤：加载模型、设置模型为评估模式、准备输入张量以及调用 `torch.onnx.export` 方法完成导出。以下是详细的说明和完整的脚本代码。 #### 1. 加载预训练的 `.pt` 模型 `.pt` 文件通常是通过 `torch.save` 保存的整个模型对象或模型状态字典。如果保存的是整个模型对象，则可以直接加载；如果是状态字典，则需要重新初始化模型并加载参数[^2]。 ```python import torch # 假设模型是以 state_dict 形式保存的 model = YourModelClass() # 替换为实际的模型类 model.load_state_dict(torch.load("your_model.pt")) # 加载模型参数 ``` #### 2. 切换模型到评估模式 在导出之前，必须确保模型处于评估模式（即调用 `model.eval()`）。这样可以禁用 Dropout 和 Batch Normalization 等训练特有的行为，使模型更适合推理阶段[^2]。 ```python model.eval() ``` #### 3. 准备示例输入张量 ONNX 导出需要一个示例输入张量来推断模型的计算图结构。该张量的形状应与模型的实际输入一致。 ```python dummy_input = torch.randn(1, 3, 224, 224) # 示例输入，假设模型接受 (N, C, H, W) 数据 ``` #### 4. 使用 `torch.onnx.export` 进行导出 `torch.onnx.export` 是 PyTorch 提供的核心函数，用于将模型导出为 ONNX 格式。以下是一些重要参数的解释： - **model**: 待导出的模型。 - **args**: 输入张量或包含多个输入的元组。 - **f**: 输出文件名或文件流。 - **export_params**: 是否导出模型参数，默认为 `True`。 - **opset_version**: 目标 ONNX 版本号，推荐使用最新稳定版（如 13 或以上）。 - **do_constant_folding**: 是否启用常量折叠优化。 - **input_names/output_names**: 定义输入/输出节点名称。 - **dynamic_axes**: 支持动态尺寸轴定义。 ```python torch.onnx.export( model, # 已经处于 eval() 状态下的模型 dummy_input, # 示例输入张量 "your_model.onnx", # 输出文件路径 export_params=True, # 包含权重参数 opset_version=13, # 使用 ONNX v13 do_constant_folding=True, # 启用常量折叠优化 input_names=["input"], # 输入节点名称 output_names=["output"], # 输出节点名称 dynamic_axes={"input": {0: "batch_size"}, # 批次维度可变 "output": {0: "batch_size"}} ) print("Model has been successfully exported to ONNX format.") ``` --- ### 完整脚本代码 以下是一个完整的 Python 脚本，展示如何从 `.pt` 文件加载模型并通过 `eval` 模式将其导出为 ONNX 格式： ```python import torch def load_and_export_model(pt_file_path, onnx_file_path): """ Load a PyTorch .pt model and export it to ONNX format. Args: pt_file_path (str): Path to the .pt file containing the saved model or its state dict. onnx_file_path (str): Output path for the generated ONNX file. """ # Step 1: Define your model class (replace with actual model definition) class YourModelClass(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) self.relu = torch.nn.ReLU() def forward(self, x): return self.relu(self.conv(x)) # Initialize the model instance model = YourModelClass() # Step 2: Load the model's state dictionary from the .pt file try: model.load_state_dict(torch.load(pt_file_path)) print(f"Successfully loaded weights from '{pt_file_path}'.") except Exception as e: raise RuntimeError(f"Failed to load model weights: {e}") # Step 3: Set the model to evaluation mode model.eval() # Step 4: Prepare an example input tensor dummy_input = torch.randn(1, 3, 224, 224) # Step 5: Export the model to ONNX format torch.onnx.export( model, dummy_input, onnx_file_path, export_params=True, opset_version=13, do_constant_folding=True, input_names=["input"], output_names=["output"], dynamic_axes={ "input": {0: "batch_size"}, "output": {0: "batch_size"} } ) print(f"Model has been successfully exported to '{onnx_file_path}'.") if __name__ == "__main__": pt_file_path = "your_model.pt" # 替换为实际的 .pt 文件路径 onnx_file_path = "your_model.onnx" # 替换为目标 ONNX 文件路径 load_and_export_model(pt_file_path, onnx_file_path) ``` --- ### 注意事项 - 如果模型中存在自定义操作或复杂控制流，可能需要额外处理才能兼容 ONNX 表达能力[^4]。 - 推荐使用最新的 PyTorch 和 ONNXRuntime 版本来获得更好的支持和性能。 - 动态批次大小的支持依赖于 `dynamic_axes` 参数配置正确。 ---