### 将PyTorch模型转换为ONNX格式并运行于CPU和GPU #### 转换PyTorch模型至ONNX格式 为了将PyTorch模型导出成ONNX格式，可以利用`torch.onnx.export()`函数。此过程涉及定义输入张量形状以及指定输出文件名。 ```python import torch import torchvision.models as models # 假设已经有一个预训练好的ResNet18模型 model = models.resnet18(pretrained=True) # 设置模型为评估模式 model.eval() # 创建一个随机输入张量作为示例输入 dummy_input = torch.randn(1, 3, 224, 224) # 导出模型到onnx文件 output_onnx = 'resnet18.onnx' torch.onnx.export(model, dummy_input, output_onnx, opset_version=11) ``` 上述代码片段展示了如何创建一个简单的ResNet18网络实例，并将其结构连同权重一起保存到了名为'resnet18.onnx'的文件中[^1]。 #### 在不同设备上加载与执行ONNX模型 一旦拥有了`.onnx`文件，则可以通过安装ONNX Runtime来实现跨平台推理支持。对于Python环境而言，可通过pip轻松获取该库： ```bash pip install onnxruntime ``` 接着，在程序里通过下面的方式实现在CPU/GPU上的推断操作: ```python import onnxruntime as ort import numpy as np def get_ort_session(onnx_model_path): providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] session = ort.InferenceSession(onnx_model_path, providers=providers) return session session = get_ort_session('resnet18.onnx') input_name = session.get_inputs()[0].name output_name = session.get_outputs()[0].name # 准备测试样本（这里假设我们已经有了处理过的图像） image_data = ... # shape (1, 3, 224, 224), dtype float32 normalized image data outputs = session.run([output_name], {input_name: image_data})[0] print(outputs.shape) # 输出预测结果维度 ``` 这段脚本首先初始化了一个能够自动选择可用硬件资源（优先考虑GPU）的会话对象；随后指定了具体的输入名称和输出名称用于后续的数据传递；最后完成了实际的前向传播计算得到最终分类得分矩阵[^2].