模型训练----将pth模型转换为onnx

Github代码

1 安装需要的环境

需要在虚拟环境中安装onnx和onnxruntime（GPU），环境和自己的cuda版本要对应上查询链接

激活环境，查看环境的cuda版本,我是cuda11.6 +cudnn8302，那就选择1.14吧

conda activate xxxx

python

import torch
# 查询cuda版本
print(torch.version.cuda)
# 查询cudnn版本
print(torch.backends.cudnn.version())

输入以下指令进行安装环境

# 安装onnx
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnx
# 安装GPU版
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnxruntime-gpu==1.14.0

2、模型转换

安装完成后，新建文件，把你训练好的模型文件导入，使用以下代码完成，我是用的resnet34训练的模型，记得修改最后的全连接层为你自己训练的类别数量，这样才能加载你训练好的权重，不然resnet34默认分类是1000

import onnx
from onnx import numpy_helper
import torch
from model import resnet34
import torch.nn as nn

pth_path = './resNet34-bird.pth'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
example = torch.randn(1,3, 224, 224).to(device)     # 1 3 224 224
print(example.dtype)

#加载模型
model = resnet34()     
in_channel = model.fc.in_features
model.fc = nn.Linear(in_channel, 14)  # 更改全连接层的输出特征数量为为自己训练的类别数
model.load_state_dict(torch.load(pth_path)) 
model = model.to(device)                            
model.eval()

# 导出模型
torch.onnx.export(model, example, r"resnet34-bird.onnx")     
model_onnx = onnx.load(r"resnet34-bird.onnx")                   # onnx加载保存的onnx模型
onnx.checker.check_model(model_onnx)                    # 检查模型是否有问题
print(onnx.helper.printable_graph(model_onnx.graph))    # 打印onnx网络

3、测试onnx模型

首先使用transforms来执行图片的加载、缩放、裁剪裁剪、转换为Tensor以及归一化等预处理步骤。然后，我们使用onnxruntime库来加载ONNX模型，并传入预处理后的图片进行推理。对于1.9版本以下的onnxruntime库需要指定在那个设备上推理模型

from PIL import Image
import numpy as np
import onnxruntime

#测试onnx模型
# Load the image
image_path = "./丹顶鹤-001.jpg"
image = Image.open(image_path)

# Preprocess the image
image = image.resize((224, 224))
image = np.array(image, dtype=np.float32)
image /= 255.0
image = np.expand_dims(image, axis=0)

# Load the ONNX model
session = onnxruntime.InferenceSession("resnet34-bird.onnx")

# Run the inference
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
output = session.run([output_name], {input_name: image})

# Get the predicted class
predicted_class = np.argmax(output)

# Print the predicted class
print("Predicted class:", predicted_class)                

可以看到最后输出结果是正确的

	点击访问博客查看更多内容