好的，下面是使用 PyTorch 来定义一个预训练过的 ResNet18 模型、加载默认模型权重、生成虚拟化输入以及将模型转换并保存为 ONNX 文件的具体步骤。最后我们还会简要说明如何验证导出的 ONNX 模型。 ### 一、安装必要的库 首先需要确保已安装PyTorch 和ONNX相关的包： ```bash pip install torch onnx onnxruntime ``` ### 二、导入所需模块及初始化ResNet18模型 ```python import torch from torchvision import models import torch.onnx as onnx_helper import numpy as np # 加载预训练的ResNet18模型并且设置为评估模式 model = models.resnet18(pretrained=True) model.eval() ``` ### 三、创建示例张量作为输入数据 为了能够顺利地完成转换过程，我们需要提供给定形状和类型的样本输入数据供模型推理: ```python x = torch.randn(1, 3, 224, 224) # batch_size=1; channels=3 (RGB); height & width of image=224x224 pixels print("Input tensor shape:", x.shape) ``` ### 四、将模型转储成ONNX格式 接下来就是核心部分——把PyTorch模型转化为ONNX文件了！ ```python output_file_name = "resnet18.onnx" torch_out = model(x) # 导出到ONNX onnx_helper.export(model, x, # 输入的数据 output_file_name,# 输出路径+名称 export_params=True, # 存储完整的参数信息 opset_version=10, # 版本号可以根据需求调整 do_constant_folding=True, # 是否执行常数折叠优化 input_names=['input'], # 给予输入节点命名 output_names=['output'] # 给予输出节点命名 ) print(f"Model has been converted to {output_file_name}") ``` ### 五、检查ONNX模型是否有效 这里我们会通过`onnxruntime`工具来进行简单的预测测试，以此来确认刚刚导出的模型能否正常工作。 ```python import onnxruntime as ort ort_session = ort.InferenceSession(output_file_name) def to_numpy(tensor): return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy() # compute ONNX Runtime output prediction ort_inputs = {'input': to_numpy(x)} ort_outputs = ort_session.run(None, ort_inputs)[0] # compare ONNX Runtime and PyTorch results np.testing.assert_allclose(to_numpy(torch_out), ort_outputs, rtol=1e-03, atol=1e-05) print("Exported model has been tested with ONNXRuntime, result looks good!") ``` 以上就是在Python环境下利用PyTorch框架定义了一个基于ImageNet数据集上预先训练完毕的标准版本ResNet18网络结构，并成功将其序列化成为ONNX形式的过程简介。 --