pth2onnx

import onnx  # 注意这里导入onnx时必须在torch导入之前，否则可能会出现segmentation fault，别人说的
import torch
from model_resnet18 import model # 这里是我的模型

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 准备模型
model_file_path = 'save_models/model-epoch10-train_acc99.3900-valid_acc98.6753.pth'  # 模型权重文件
checkpoint = torch.load(model_file_path)
model.load_state_dict(checkpoint)  # 模型加载训练好权重
model.to(device)
# model.eval()  # 评估模式,非训练模式,该模式下，batchNorm层、dropout层等用于优化训练而添加的网络层会被关闭,和下面export()中的traing=False重复了？

# 开始转换
dummy_input = torch.randn(10, 3, 224, 224, device=device)  # 虚拟输入，用于确定输入尺寸和推理计算图每个节点的尺寸
output_model_name = "classifier_3x224x224_batchsize_10.onnx"
input_names = ["input"]  # 输入节点的名称，这里也可以给一个list，list中名称分别对应每一层可学习的参数，便于后续查询
output_names = ["output"]  # 输出节点的名称

output = torch.onnx.export(model,
                           dummy_input,
                           output_model_name,  # 输出文件的名称
                           export_params=True, # 如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False
                           verbose=False,  # 是否以字符串的形式显示计算图
                           #training=False, # 在训练模式下导出模型。目前，ONNX导出的模型只是为了做推断，所以你通常不需要将其设置为True.  后面有一次报错 ，注释了'mode' should be a torch.onnx.TrainingMode enum, but got '<class 'bool'>'.
                           input_names=input_names,
                           output_names=output_names,
                           do_constant_folding=True,  # 是否压缩常量
                           #dynamic_axes={"input": {0: "batch_size", 2: "h"}, "output": {0: "batch_size"}, },  # 设置动态维度，此处指明input节点的第0维度可变，命名为batch_size
                           #dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}, }
                           )




# 加载onnx模型并检查,输出计算图
net = onnx.load(output_model_name)  # 加载onnx计算图
onnx.checker.check_model(net)  # 检查模型格式是否完整及正确
graph = onnx.helper.printable_graph(net.graph)  # 输出onnx的计算图
print(graph)



# 浏览器中可视化
import netron

netron.start("classifier_1x224x224_20220707.onnx")  # 默认会打开8080端口，会直接跳转到浏览器中，可视化

 查看输出维度

import onnxruntime
import numpy as np

session = onnxruntime.InferenceSession("classifier_1x224x224_20220707.onnx")  # 创建一个运行session，类似于tensorflow
out_r = session.run(None, {"input": np.random.rand(16, 3, 224, 224).astype('float32')})  # 模型运行，注意这里的输入必须是numpy类型
print(len(out_r))
print(out_r[0].shape)  # 打印输出shape

 

yolo onnx模型验证

 

ppn onnx模型验证

# 读取并处理图片
inp = cv2.imread(img_path) 
inp = scale_img(inp, (224, 224))  # inp.shape: (384, 384, 3)
inp = inp[:, :, ::-1] / 255  # BGR convert to RGB,相当于把每一行反转了
inp = np.expand_dims(inp, axis=0)  # inp.shape: (1, 384, 384, 3)
inp_ = (inp - np.array([0.485, 0.456, 0.406])) / np.array([0.229, 0.224, 0.225])  # inp_.shape: (1, 384, 384, 3)
# [b h w c]-->[b c h w]
inp_ = torch.from_numpy(inp_.transpose((0, 3, 1, 2)).astype(np.float32))
inp_ = inp_.to(device)  # inp_.shape:(1,3,384,384)

# 加载onnx模型
session = onnxruntime.InferenceSession("d:/05jule/jule_data_224x224_5x5_430_1.73.onnx", None)
output = session.run([], {"input": inp_.cpu().numpy()}) # 这里传入的是numpy格式,返回值是一个list，所以后面的代码为output[0]
output_numpy = np.array(output[0])
output = torch.from_numpy(output_numpy).to(device)

# 后面就是后处理的代码，查看结果