还在为模型加速推理发愁吗？不如看看这篇吧。手把手教你把pytorch模型转化为TensorRT，加速推理

现在能在网上找到很多很多的学习资源，有免费的也有收费的，当我拿到1套比较全的学习资源之前，我并没着急去看第1节，我而是去审视这套资源是否值得学习，有时候也会去问一些学长的意见，如果可以之后，我会对这套学习资源做1个学习计划，我的学习计划主要包括规划图和学习进度表。

分享给大家这份我薅到的免费视频资料，质量还不错，大家可以跟着学习

网上学习资料一大堆，但如果学到的知识不成体系，遇到问题时只是浅尝辄止，不再深入研究，那么很难做到真正的技术提升。

需要这份系统化学习资料的朋友，可以戳这里获取

一个人可以走的很快，但一群人才能走的更远！不论你是正从事IT行业的老鸟或是对IT行业感兴趣的新人，都欢迎加入我们的的圈子（技术交流、学习资源、职场吐槽、大厂内推、面试辅导），让我们一起学习成长！

• 摘要

• 环境配置

• ONNX

• 模型转化

• • 第一种推理写法

• 第二种推理写法

• TensorRT

• • 模型转化

• 动态推理

• 静态推理

摘要

=============================================================

最近，学习了一些模型转化和加速推理的知识，本文是对学习成果的总结。

对模型的转化，本文实现了pytorch模型转onnx模型和onnx转TensorRT，在转为TensorRT模型的过程中，实现了模型单精度的压缩。

对于加速推理，本文实现GPU环境下的onnxruntime推理、TensorRT动态推理和TensorRT静态推理。

希望本文能帮助大家。

环境配置

===============================================================

CUDA版本：11.3.1

cuDNN版本：8.2.1

TensorRT版本：8.0.3.4

显卡：1650

pytorch：1.10.2

模型的转化和推理对版本有要求，如果版本对应不上很可能出现千奇百怪的问题，所以我把我的版本信息列出来给大家做参考。

ONNX

===============================================================

ONNX，全称：Open Neural Network Exchange（ONNX，开放神经网络交换），是一个用于表示深度学习模型的标准，可使模型在不同框架之间进行转移。

ONNX是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如Pytorch, MXNet）可以采用相同格式存储模型数据并交互。 ONNX的规范及代码主要由微软，亚马逊 ，Facebook 和 IBM 等公司共同开发，以开放源代码的方式托管在Github上。目前官方支持加载ONNX模型并进行推理的深度学习框架有： Caffe2, PyTorch, MXNet，ML.NET，TensorRT 和 Microsoft CNTK，并且 TensorFlow 也非官方的支持ONNX。—维基百科

onnx模型可以看作是模型转化之间的中间模型，同时也是支持做推理的。一般来说，onnx的推理速度要比pytorch快上一倍。

模型转化

===============================================================

onnx模型转换和推理需要安装Python包，如下：

pip install onnx

pip install onnxruntime-gpu

新建模型转换脚本pytorch2onnx.py。

import torch

from torch.autograd import Variable

import onnx

import netron

print(torch.version)

input_name = [‘input’]

output_name = [‘output’]

input = Variable(torch.randn(1, 3, 224, 224)).cuda()

model = torch.load(‘model.pth’, map_location=“cuda:0”)

torch.onnx.export(model, input, ‘model_onnx.onnx’,opset_version=13, input_names=input_name, output_names=output_name, verbose=True)

模型可视化

netron.start(‘model_onnx.onnx’)

导入需要的包。

打印pytorch版本。

定义input_name和output_name变量。

定义输入格式。

加载pytorch模型。

导出onnx模型，这里注意一下参数opset_version在8.X版本中设置为13，在7.X版本中设置为12。

yolov5中这么写的。

if trt.version[0] == ‘7’: # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012

grid = model.model[-1].anchor_grid

model.model[-1].anchor_grid = [a[…, :1, :1, :] for a in grid]

export_onnx(model, im, file, 12, train, False, simplify) # opset 12

model.model[-1].anchor_grid = grid

else: # TensorRT >= 8

check_version(trt.version, ‘8.0.0’, hard=True) # require tensorrt>=8.0.0

export_onnx(model, im, file, 13, train, False, simplify) # opset 13

查看转化后的模型，如下图：

推理的写法有两种，一种直接写，另一种将其封装为通用的推理类。

第一种推理写法

---

先看第一种写法，新建test_onnx.py，插入下面的代码：

import os, sys

import time

sys.path.append(os.getcwd())

import onnxruntime

import numpy as np

import torchvision.transforms as transforms

from PIL import Image

导入包

def get_test_transform():

return transforms.Compose([

transforms.Resize([224, 224]),

transforms.ToTensor(),

transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),

])

image = Image.open(‘11.jpg’) # 289

img = get_test_transform()(image)

img = img.unsqueeze_(0) # -> NCHW, 1,3,224,224

print(“input img mean {} and std {}”.format(img.mean(), img.std()))

img = np.array(img)

定义get_test_transform函数，实现图像的归一化和resize。

读取图像。

对图像做resize和归一化。

增加一维batchsize。

将图片转为数组。

onnx_model_path = “model_onnx.onnx”

##onnx测试

session = onnxruntime.InferenceSession(onnx_model_path,providers=[‘TensorrtExecutionProvider’, ‘CUDAExecutionProvider’, ‘CPUExecutionProvider’])

#compute ONNX Runtime output prediction

inputs = {session.get_inputs()[0].name: img}

time3=time.time()

outs = session.run(None, inputs)[0]

y_pred_binary = np.argmax(outs, axis=1)

print(“onnx prediction”, y_pred_binary[0])

time4=time.time()

print(time4-time3)

定义onnx_model_path模型的路径。

加载onnx模型。

定义输入。

执行推理。

获取预测结果。

到这里第一种写法就完成了，是不是很简单，接下来看第二种写法。

第二种推理写法

---

新建onnx.py脚本，加入以下代码：

import onnxruntime

class ONNXModel():

def init(self, onnx_path):

“”"

:param onnx_path:

“”"

self.onnx_session = onnxruntime.InferenceSession(onnx_path,providers=[‘TensorrtExecutionProvider’, ‘CUDAExecutionProvider’, ‘CPUExecutionProvider’])

self.input_name = self.get_input_name(self.onnx_session)

self.output_name = self.get_output_name(self.onnx_session)

print(“input_name:{}”.format(self.input_name))

print(“output_name:{}”.format(self.output_name))

def get_output_name(self, onnx_session):

“”"

output_name = onnx_session.get_outputs()[0].name

:param onnx_session:

:return:

“”"

output_name = []

for node in onnx_session.get_outputs():

output_name.append(node.name)

return output_name

def get_input_name(self, onnx_session):

“”"

input_name = onnx_session.get_inputs()[0].name

:param onnx_session:

:return:

“”"

input_name = []

for node in onnx_session.get_inputs():

input_name.append(node.name)

return input_name

def get_input_feed(self, input_name, image_numpy):

“”"

input_feed={self.input_name: image_numpy}

:param input_name:

:param image_numpy:

:return:

“”"

input_feed = {}

for name in input_name:

input_feed[name] = image_numpy

return input_feed

def forward(self, image_numpy):

输入数据的类型必须与模型一致,以下三种写法都是可以的

scores, boxes = self.onnx_session.run(None, {self.input_name: image_numpy})

scores, boxes = self.onnx_session.run(self.output_name, input_feed={self.input_name: iimage_numpy})

input_feed = self.get_input_feed(self.input_name, image_numpy)

scores = self.onnx_session.run(self.output_name, input_feed=input_feed)

return scores

调用onnx.py实现推理，新建test_onnx1.py插入代码：

import os, sys

sys.path.append(os.getcwd())

import numpy as np

import torchvision.transforms as transforms

from PIL import Image

from models.onnx import ONNXModel

def get_test_transform():

return transforms.Compose([

transforms.Resize([224, 224]),

transforms.ToTensor(),

transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),

])

image = Image.open(‘11.jpg’) # 289

img = get_test_transform()(image)

img = img.unsqueeze_(0) # -> NCHW, 1,3,224,224

print(“input img mean {} and std {}”.format(img.mean(), img.std()))

img = np.array(img)

onnx_model_path = “model_onnx.onnx”

model1 = ONNXModel(onnx_model_path)

out = model1.forward(img)

y_pred_binary = np.argmax(out[0], axis=1)

print(“onnx prediction1”, y_pred_binary[0])

输出结果如下：

TensorRT

===================================================================

TensorRT是英伟达推出的一个高性能的深度学习推理（Inference）优化器，可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架，将TensorRT和NVIDIA的GPU结合起来，能在几乎所有的框架中进行快速和高效的部署推理。

TensorRT 是一个C++库，从 TensorRT 3 开始提供C++ API和Python API，主要用来针对 NVIDIA GPU进行 高性能推理（Inference）加速。

TensorRT的安装可以参考我以前的文章：

https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/120360288。

本次用的8.X版本的，安装方式一样。

本文实现Python版本的 TensorRT 推理加速,需要安装tensorrt包文件。这个文件不能直接通过pip下载，我在下载的TensorRT安装包里，不过我下载的8.0.3.4版本中并没有，在8.2.1.8的版本中存在这个包文件。

所以我安装了8.2.1.8中的whl文件。

安装方式，进入模型所在的目录，执行：

pip install tensorrt-8.2.1.8-cp39-none-win_amd64.whl

学好 Python 不论是就业还是做副业赚钱都不错，但要学会 Python 还是要有一个学习规划。最后大家分享一份全套的 Python 学习资料，给那些想学习 Python 的小伙伴们一点帮助！

一、Python所有方向的学习路线

Python所有方向路线就是把Python常用的技术点做整理，形成各个领域的知识点汇总，它的用处就在于，你可以按照上面的知识点去找对应的学习资源，保证自己学得较为全面。

二、学习软件

工欲善其事必先利其器。学习Python常用的开发软件都在这里了，给大家节省了很多时间。

三、全套PDF电子书

书籍的好处就在于权威和体系健全，刚开始学习的时候你可以只看视频或者听某个人讲课，但等你学完之后，你觉得你掌握了，这时候建议还是得去看一下书籍，看权威技术书籍也是每个程序员必经之路。

四、入门学习视频

我们在看视频学习的时候，不能光动眼动脑不动手，比较科学的学习方法是在理解之后运用它们，这时候练手项目就很适合了。

五、实战案例

光学理论是没用的，要学会跟着一起敲，要动手实操，才能将自己的所学运用到实际当中去，这时候可以搞点实战案例来学习。

六、面试资料

我们学习Python必然是为了找到高薪的工作，下面这些面试题是来自阿里、腾讯、字节等一线互联网大厂最新的面试资料，并且有阿里大佬给出了权威的解答，刷完这一套面试资料相信大家都能找到满意的工作。

网上学习资料一大堆，但如果学到的知识不成体系，遇到问题时只是浅尝辄止，不再深入研究，那么很难做到真正的技术提升。

需要这份系统化学习资料的朋友，可以戳这里获取

一个人可以走的很快，但一群人才能走的更远！不论你是正从事IT行业的老鸟或是对IT行业感兴趣的新人，都欢迎加入我们的的圈子（技术交流、学习资源、职场吐槽、大厂内推、面试辅导），让我们一起学习成长！