1. pytorch模型转换到onnx模型
2.运行onnx模型
3.比对onnx模型和pytorch模型的输出结果
我这里重点是第一点和第二点,第三部分 比较容易
首先你要安装 依赖库:onnx 和 onnxruntime,
- pip install onnx
- pip install onnxruntime 进行安装
也可以使用清华源镜像文件安装 速度会快些。
开始:
1. pytorch模型转换到onnx模型
pytorch 转 onnx 仅仅需要一个函数 torch.onnx.export
torch.onnx.export(model, args, path, export_params, verbose, input_names, output_names, do_constant_folding, dynamic_axes, opset_version)
参数说明:
转化代码:参考1:
- import torch
- import torch.nn
- import onnx
-
- model = torch.load('best.pt')
- model.eval()
-
- input_names = ['input']
- output_names = ['output']
-
- x = torch.randn(1,3,32,32,requires_grad=True)
-
- torch.onnx.export(model, x, 'best.onnx', input_names=input_names, output_names=output_names, verbose='True')
参考2:PlainC3AENetCBAM 是网络模型,如果你没有自己的网络模型,可能成功不了
- import io
- import torch
- import torch.onnx
- from models.C3AEModel import PlainC3AENetCBAM
-
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
- def test():
- model = PlainC3AENetCBAM()
-
- pthfile = r'/home/joy/Projects/models/emotion/PlainC3AENet.pth'
- loaded_model = torch.load(pthfile, map_location='cpu')
- # try:
- # loaded_model.eval()
- # except AttributeError as error:
- # print(error)
-
- model.load_state_dict(loaded_model['state_dict'])
- # model = model.to(device)
-
- #data type nchw
- dummy_input1 = torch.randn(1, 3, 64, 64)
- # dummy_input2 = torch.randn(1, 3, 64, 64)
- # dummy_input3 = torch.randn(1, 3, 64, 64)
- input_names = [ "actual_input_1"]
- output_names = [ "output1" ]
- # torch.onnx.export(model, (dummy_input1, dummy_input2, dummy_input3), "C3AE.onnx", verbose=True, input_names=input_names, output_names=output_names)
- torch.onnx.export(model, dummy_input1, "C3AE_emotion.onnx", verbose=True, input_names=input_names, output_names=output_names)
-
- if __name__ == "__main__":
- test()
直接将PlainC3AENetCBAM替换成需要转换的模型,然后修改pthfile,输入和onnx模型名字然后执行即可。
注意:上面代码中注释的dummy_input2,dummy_input3,torch.onnx.export对应的是多个输入的例子。
在转换过程中遇到的问题汇总
RuntimeError: Failed to export an ONNX attribute, since it's not constant, please try to make things (e.g., kernel size) static if possible
在转换过程中遇到RuntimeError: Failed to export an ONNX attribute, since it's not constant, please try to make things (e.g., kernel size) static if possible的错误。
我成功的案例,我直接把我训练的网络贴上,成功转换,没有from ** import 模型名词这么委婉,合法,我的比较粗暴
- import torch
- import torch.nn
- import onnx
- from torchvision import transforms
- import torch.nn as nn
- from torch.nn import Sequential
-
- # 添加模型
-
- # 设置数据转换方式
- preprocess_transform = transforms.Compose([
- transforms.ToTensor(), # 把数据转换为张量(Tensor)
- transforms.Normalize( # 标准化,即使数据服从期望值为 0,标准差为 1 的正态分布
- mean=[0.5, ], # 期望
- std=[0.5, ] # 标准差
- )
- ])
-
- class CNN(nn.Module): # 从父类 nn.Module 继承
- def __init__(self): # 相当于 C++ 的构造函数
- # super() 函数是用于调用父类(超类)的一个方法,是用来解决多重继承问题的
- super(CNN, self).__init__()
-
- # 第一层卷积层。Sequential(意为序列) 括号内表示要进行的操作
- self.conv1 = Sequential(
- nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=1, padding=1),
- nn.BatchNorm2d(64),
- nn.ReLU(),
- nn.MaxPool2d(kernel_size=2, stride=2)
- )
-
- # 第二卷积层
- self.conv2 = Sequential(
- nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),
- nn.BatchNorm2d(128),
- nn.ReLU(),
- nn.MaxPool2d(kernel_size=2, stride=2)
- )
-
- # 全连接层(Dense,密集连接层)
- self.dense = Sequential(
- nn.Linear(7 * 7 * 128, 1024),
- nn.ReLU(),
- nn.Dropout(p=0.5),
- nn.Linear(1024, 10)
- )
-
- def forward(self, x): # 正向传播
- x1 = self.conv1(x)
- x2 = self.conv2(x1)
- x = x2.view(-1, 7 * 7 * 128)
- x = self.dense(x)
- return x
-
- # 训练
- # 训练和参数优化
-
- # 定义求导函数
- def get_Variable(x):
- x = torch.autograd.Variable(x) # Pytorch 的自动求导
- # 判断是否有可用的 GPU
- return x.cuda() if torch.cuda.is_available() else x
-
-
- # 判断是否GPU
- device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
- # device1 = torch.device('cpu')
- # 定义网络
- model = CNN()
-
-
-
-
- loaded_model = torch.load('save_model/model.pth', map_location='cuda:0')
- model.load_state_dict(loaded_model)
- model.eval()
-
- input_names = ['input']
- output_names = ['output']
-
- # x = torch.randn(1,3,32,32,requires_grad=True)
- x = torch.randn(1, 1, 28, 28, requires_grad=True) # 这个要与你的训练模型网络输入一致。我的是黑白图像
-
- torch.onnx.export(model, x, 'save_model/model.onnx', input_names=input_names, output_names=output_names, verbose='True')
-
前提是你要准备好*.pth模型保持文件
输出结果:
- graph(%input : Float(1, 1, 28, 28, strides=[784, 784, 28, 1], requires_grad=1, device=cpu),
- %dense.0.weight : Float(1024, 6272, strides=[6272, 1], requires_grad=1, device=cpu),
- %dense.0.bias : Float(1024, strides=[1], requires_grad=1, device=cpu),
- %dense.3.weight : Float(10, 1024, strides=[1024, 1], requires_grad=1, device=cpu),
- %dense.3.bias : Float(10, strides=[1], requires_grad=1, device=cpu),
- %33 : Float(64, 1, 3, 3, strides=[9, 9, 3, 1], requires_grad=0, device=cpu),
- %34 : Float(64, strides=[1], requires_grad=0, device=cpu),
- %36 : Float(128, 64, 3, 3, strides=[576, 9, 3, 1], requires_grad=0, device=cpu),
- %37 : Float(128, strides=[1], requires_grad=0, device=cpu)):
- %input.4 : Float(1, 64, 28, 28, strides=[50176, 784, 28, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%input, %33, %34) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\modules\conv.py:443:0
- %21 : Float(1, 64, 28, 28, strides=[50176, 784, 28, 1], requires_grad=1, device=cpu) = onnx::Relu(%input.4) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:1442:0
- %input.8 : Float(1, 64, 14, 14, strides=[12544, 196, 14, 1], requires_grad=1, device=cpu) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%21) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:797:0
- %input.16 : Float(1, 128, 14, 14, strides=[25088, 196, 14, 1], requires_grad=1, device=cpu) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%input.8, %36, %37) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\modules\conv.py:443:0
- %25 : Float(1, 128, 14, 14, strides=[25088, 196, 14, 1], requires_grad=1, device=cpu) = onnx::Relu(%input.16) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:1442:0
- %26 : Float(1, 128, 7, 7, strides=[6272, 49, 7, 1], requires_grad=1, device=cpu) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%25) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:797:0
- %27 : Long(2, strides=[1], device=cpu) = onnx::Constant[value= -1 6272 [ CPULongType{2} ]]() # E:/paddle_project/Pytorch_Imag_Classify/zifu_fenlei/CNN/pt模型转onnx模型.py:51:0
- %28 : Float(1, 6272, strides=[6272, 1], requires_grad=1, device=cpu) = onnx::Reshape(%26, %27) # E:/paddle_project/Pytorch_Imag_Classify/zifu_fenlei/CNN/pt模型转onnx模型.py:51:0
- %input.20 : Float(1, 1024, strides=[1024, 1], requires_grad=1, device=cpu) = onnx::Gemm[alpha=1., beta=1., transB=1](%28, %dense.0.weight, %dense.0.bias) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\modules\linear.py:103:0
- %input.24 : Float(1, 1024, strides=[1024, 1], requires_grad=1, device=cpu) = onnx::Relu(%input.20) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\functional.py:1442:0
- %output : Float(1, 10, strides=[10, 1], requires_grad=1, device=cpu) = onnx::Gemm[alpha=1., beta=1., transB=1](%input.24, %dense.3.weight, %dense.3.bias) # D:\ProgramData\Anaconda3\envs\openmmlab\lib\site-packages\torch\nn\modules\linear.py:103:0
- return (%output)
输出结果的device 是CPU,模型加载的时候是GPU。这就是转换的意义吧
2.运行onnx模型
- import onnx
- import onnxruntime as ort
-
- model = onnx.load('best.onnx')
- onnx.checker.check_model(model)
-
- session = ort.InferenceSession('best.onnx')
- x=np.random.randn(1,3,32,32).astype(np.float32) # 注意输入type一定要np.float32!!!!!
- # x= torch.randn(batch_size,chancel,h,w)
-
-
- outputs = session.run(None,input = { 'input' : x })
参考: