ONNX转换及如何使用的代码

我们保存模型的过程中，一般而言是保存state_dict格式

torch.save(model.state_dict(), 'model_xx_classfiy.pth')

保存的state_dict格式的pth文件中仅仅含有的是模型的参数信息，而不具备模型的结构信息，因此其不具备直接使用的功能，在我们使用的过程中（load后）需要将参数赋予模型（简单来说就是创建一个模型，将参数赋予这个模型，然后将其应用于诸如测试、转模型的工作）

如何进行转onnx操作

代码如下：

import torch

from model import Net

model_test = Net()

model_statedict = torch.load("model_figure_classfiy_e500.pth",map_location=lambda storage,loc:storage)   #导入Gpu训练模型，导入为cpu格式

model_test.load_state_dict(model_statedict)  #将参数放入model_test中

model_test.eval()  # 测试，看是否报错

#下面开始转模型，cpu格式下

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

dummy_input = torch.randn(1, 3, 64, 64,device=device)

input_names = ["input"]

output_names = ["output"]

torch.onnx.export(model_test, dummy_input, "model_.onnx", opset_version=9, verbose=False )

这样就可以转型成功了，但是要注意输入这边的变化，要根据你的模型输入来适当调整。

如何使用转换后的onnx模型

import cv2 as cv
import numpy as np

def img_process(image):
    mean = np.array([0.5,0.5,0.5],dtype=np.float32).reshape(1,1,3)
    std = np.array([0.5,0.5,0.5],dtype=np.float32).reshape(1,1,3)
    new_img = ((image/255. -mean)/std).astype(np.float32)
    return new_img

img = cv.imread("figure_1.jpg")
img_t = cv.resize(img,(64,64))    #将图片改为模型适用的尺寸
img_t = img_process(img_t)

layerNames = ["hm"]   # 这里的输出的名称应该于前面的转模型时候定义的一致
blob=cv.dnn.blobFromImage(img_t,scalefactor=1.0,swapRB=True,crop=False)  # 将image转化为 1x3x64x64 格式输入模型中
net = cv.dnn.readNetFromONNX("model_.onnx")
net.setInput(blob)
outs = net.forward(layerNames)
print(outs)

注意在输入模型的图像要进行标准化预处理使其与onnx模型输入保持一致，否则会报错。

这一部分只要用于模型推理，可以根据这个去验证我们模型转换后，模型是否转换正确，是否很好的保存了模型的参数。