### 如何使用PyTorch训练模型并保存为PT文件 在PyTorch中，可以通过`torch.save()`函数将训练好的模型保存为`.pt`文件。以下是一个简单的代码示例： ```python import torch import torch.nn as nn # 定义一个简单的神经网络模型 class SimpleModel(nn.Module): def __init__(self): super(SimpleModel, self).__init__() self.fc = nn.Linear(10, 1) def forward(self, x): return self.fc(x) # 创建模型实例 model = SimpleModel() # 假设模型已经训练完成 # 保存模型参数到pt文件 torch.save(model.state_dict(), "model.pt") # 仅保存模型参数 ``` 通过上述代码，可以将模型的参数保存到`model.pt`文件中[^3]。 --- ### 将PT文件转换为ONNX格式 为了将保存的`.pt`文件转换为ONNX格式，需要加载模型并调用`torch.onnx.export()`函数。以下是具体实现： ```python # 加载模型参数 model.load_state_dict(torch.load("model.pt")) model.eval() # 设置模型为评估模式 # 定义输入张量的形状 dummy_input = torch.randn(1, 10) # 示例输入，形状为 (batch_size, input_dim) # 导出为ONNX格式 torch.onnx.export( model, # 要导出的模型 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"}} # 动态轴定义 ) ``` 通过上述代码，可以将`.pt`文件成功转换为ONNX格式文件`model.onnx`[^4]。 --- ### 使用onnx-simplifier简化ONNX模型 为了简化生成的ONNX模型，可以使用`onnx-simplifier`库。以下是具体操作步骤： 1. **安装onnx-simplifier**：如果尚未安装该库，可以通过以下命令安装： ```bash pip install onnx-simplifier ``` 2. **简化ONNX模型**：使用`onnxsim.simplify()`函数对模型进行简化： ```python import onnx from onnxsim import simplify # 加载ONNX模型 model = onnx.load("model.onnx") # 简化模型 model_simp, check = simplify(model) # 验证简化是否成功 assert check, "Simplified ONNX model could not be validated" # 保存简化后的模型 onnx.save(model_simp, "model_simplified.onnx") ``` 通过上述代码，可以将原始的ONNX模型简化并保存为`model_simplified.onnx`[^5]。 --- ### 注意事项 - 在将PyTorch模型导出为ONNX时，确保模型支持ONNX操作集。某些自定义操作可能无法直接转换。 - `onnx-simplifier`可能会移除冗余操作或优化计算图，但需验证简化后的模型功能是否保持一致。 - 如果模型包含动态输入大小，务必在`torch.onnx.export()`中正确设置`dynamic_axes`参数。 ---