前言

休闲大半个月，终于又要开始认真工作了，果然游戏害人（淦），废话不多说，现在开始工作，我在转模型的时候遇到总总问题，先是转化不成功（前面有文章说过这个问题），后面转成功了但是我前几天想用那个模型的时候发现，转化的模型有大问题，首先是数值上和pth文件不一致，其次甚至没有BatchNormalization模块，具体看下图（netron查看onnx模型）

这种问题可能是pytorch版本问题，建议用比较高的版本，下面给出转模型的全过程以及解决上述问题的方法。

一、pth转onnx模型

首先通过一个简单的例子来介绍转模型的过程。下面是我搭建的一个简单的用作分类的卷积网络（文件名 model.py）：

import torch.nn as nn
import torch.nn.functional as F# 定义卷积神经网络结构
class Net(nn.Module):def __init__(self):super(Net, self).__init__()# 卷积层 (32x32x3的图像)self.conv1 = nn.Conv2d(3, 16, 3,stride=1, padding=1)# 卷积层(16x16x16)self.conv2 = nn.Conv2d(16, 32, 3, stride=2,padding=1)# 卷积层(8x8x32)self.conv3 = nn.Conv2d(32, 64, 3,stride=1, padding=1)# 最大池化层self.pool = nn.MaxPool2d(2, 2)# linear layer (64 * 4 * 4 -> 500)self.fc1 = nn.Linear(64 * 4 * 4, 500)# linear layer (500 -> 10)self.fc2 = nn.Linear(500, 10)# dropout层 (p=0.3)self.dropout = nn.Dropout(0.3)def forward(self, x):# add sequence of convolutional and max pooling layersx = self.pool(F.relu(self.conv1(x)))x = self.pool(F.relu(self.conv2(x)))x = self.pool(F.relu(self.conv3(x)))# flatten image input#print(x.shape)x = x.view(-1, 64 * 4 * 4)# add dropout layerx = self.dropout(x)# add 1st hidden layer, with relu activation functionx = F.relu(self.fc1(x))# add dropout layerx = self.dropout(x)# add 2nd hidden layer, with relu activation functionx = self.fc2(x)return x

还是比较简单的，想必大家都能轻松看懂，当然懂不懂没关系，只是想让大家知道，转模型的过程中必须要import模型文件。
其次在我们保存模型的过程中，一般而言是保存state_dict格式（非state_dict格式在应用过程中会出现各种问题，pytorch官网也推荐使用state_dict格式），具体保存的函数可以见如下（仅供参考）：

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

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

import torch
from model import Netmodel_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, output_names=["hm"])

上述代码可以成功将模型转化成功，但是可能会报一个warning，后面我用opencv的dnn模块调用此onnx模型，发现可以成功。

import cv2 as cv
import numpy as npdef 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_imgimg = cv.imread("figure_1.jpg")
img_t = cv.resize(img,(64,64))    #将图片改为模型适用的尺寸
img_t = img_process(img_t)
#img_t = np.transpose(img_t,[2,0,1])
#img_t = img_t[np.newaxis,:]   #扩展一个新维度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)

结果如下：

并且我输入的图片如下：

对应的输出中下标为6的值最大为11.41577 ，分类正确。
后续我会将整个小项目的资源打包汇总一下，希望能对你有用。

其实看到这，你应该对pth模型转onnx模型有一个比较清晰的了解了，其实归根结底就是一个torch.onnx.export()函数的应用，这是此函数的一些参数介绍，希望能够帮助你更加清晰的了解：

torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None)
参数介绍：model (torch.nn.Module) – 要导出的模型.args (tuple of arguments) – 模型的输入, 任何非Tensor参数都将硬编码到导出的模型中；任何Tensor参数都将成为导出的模型的输入，并按照他们在args中出现的顺序输入。因为export运行模型，所以我们需要提供一个输入张量x。只要是正确的类型和大小，其中的值就可以是随机的。请注意，除非指定为动态轴，否则输入尺寸将在导出的ONNX图形中固定为所有输入尺寸。在此示例中，我们使用输入batch_size 1导出模型，但随后dynamic_axes 在torch.onnx.export()。因此，导出的模型将接受大小为[batch_size，3、100、100]的输入，其中batch_size可以是可变的。f - 保存后的onnx文件名export_params (bool, default True) – 如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 Falseverbose (bool, default False) - 如果指定，我们将打印出一个导出轨迹的调试描述training (bool, default False) - 在训练模式下导出模型。目前，ONNX导出的模型只是为了做推断，所以你通常不需要将其设置为Trueinput_names (list of strings, default empty list) – 按顺序分配名称到图中的输入节点output_names (list of strings, default empty list) –按顺序分配名称到图中的输出节点

以上就是整个模型的转化过程了，下面是我在转一些比复杂的模型的时候遇到的一些问题。如果遇到相同问题的可以了解一下。

模型转换：pth转onnx相关推荐

yolov5模型转换（pt=＞onnx=＞rknn）和板端验证测试
测试环境说明: (1)由于模型转换工具需要onnx版本和rknn的tool工具需要的版本相互矛盾需要创建量开发环境,当前测试转换的模型是yolov5_v5.0的模型 (2)由于在搭建开发环境时还存在部 ...
Atlas深度学习模型转换及运行
转换方法有两个主流的方法 onnx->Caffe Model->om(最主流) onnx->om 使用第一种方法看似麻烦,多了一步,但其实这种转换方式的好处有很多可以通过修改pr ...
【1】AI模型转换综述
前言当用户基于各种原因学习并使用了一种框架的时候,常常会发现应用或者再训练的场景改变了,比如用户用 Caffe 训练好了一个图像识别的模型,但是生产环境是使用 TensorFlow 做预测.再比如某 ...
将训练好的pytorch模型的pth文件转换成onnx模型（亲测成功）
将训练好的pytorch模型的pth文件转换成onnx模型(亲测成功) 模型转换声明:本文原创,未经许可严禁转载,原文地址https://blog.csdn.net/hutao1030813002/ ...
【地平线开发板模型转换】将pytorch生成的onnx模型转换成.bin模型
文章目录 1 获取onnx模型 2 启动docker容器 3 onnx模型检查 3.1 为什么要检查? 3.2 如何操作 4 图像数据预处理 4.1 一些问题的思考 4.2 图片挑选与放置 4.2 使 ...
pytorch模型转ONNX转TensorRT，模型转换和推理部署
一.pth模型转ONNX import os import sys import torch import numpy as npfrom feat.model import ResNet # 导入自 ...
模型转换：pytorch模型转onnx, onnx转tensorflow, tensorflow转tflite
文章目录软件版本: pytorch模型转onnx onnx模型转tensorflow tensorflow模型转tflite 软件版本: tensorflow 2.3.1 pytorch 1.6.0 ...
yolo模型转换：pytorch -＞ onnx -＞ caffe
第一步:pytorch转onnx(pytorch版,yolov3-9.0开始提供脚本export.py) (1)设置onnx算子版本(按需) 修改代码: torch.onnx.export(model ...
MMDetection2.17-权重模型转推理模型(pth转onnx)详细步骤及前向推理(Win10、Linux均适用)
权重模型转推理模型的意义? 换后的ONNX模型注意要点: 转换流程及完整代码: 前向推理及完整代码: 权重模型转推理模型的意义? 方便部署:转为onnx格式的模型后,就可以不需要依赖mmdetecti ...
关于在SNPE平台上进行ONNX模型转换DLC模型
Onnx模型转化DLC模型简介在snpe平台上,将onnx模型转换为dlc模型目录 snpe平台介绍 snpe平台与onnx配置 onnx模型转换dlc 模型量化关于1.38版本SNPE部署时 ...

模型转换：pth转onnx

文章目录

前言

一、pth转onnx模型

模型转换：pth转onnx相关推荐

最新文章

热门文章