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

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

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

二、Python必备开发工具

工具都帮大家整理好了，安装就可直接上手！

三、最新Python学习笔记

当我学到一定基础，有自己的理解能力的时候，会去阅读一些前辈整理的书籍或者手写的笔记资料，这些笔记详细记载了他们对一些技术点的理解，这些理解是比较独到，可以学到不一样的思路。

四、Python视频合集

观看全面零基础学习视频，看视频学习是最快捷也是最有效果的方式，跟着视频中老师的思路，从基础到深入，还是很容易入门的。

五、实战案例

纸上得来终觉浅，要学会跟着视频一起敲，要动手实操，才能将自己的所学运用到实际当中去，这时候可以搞点实战案例来学习。

六、面试宝典

简历模板

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

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

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

增加一维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

模型推理用到了pycuda，执行安装命令：

pip install pycuda

模型转化

---

将onnx模型转为TensorRT 模型，新建onnx2trt.py，插入代码：

import tensorrt as trt

def build_engine(onnx_file_path,engine_file_path,half=False):

“”“Takes an ONNX file and creates a TensorRT engine to run inference with”“”

logger = trt.Logger(trt.Logger.INFO)

builder = trt.Builder(logger)

config = builder.create_builder_config()

config.max_workspace_size = 4 * 1 << 30

flag = (1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

network = builder.create_network(flag)

parser = trt.OnnxParser(network, logger)

if not parser.parse_from_file(str(onnx_file_path)):

raise RuntimeError(f’failed to load ONNX file: {onnx_file_path}')

half &= builder.platform_has_fast_fp16

if half:

config.set_flag(trt.BuilderFlag.FP16)

with builder.build_engine(network, config) as engine, open(engine_file_path, ‘wb’) as t:

t.write(engine.serialize())

return engine_file_path

if name ==“main”:

onnx_path1 = ‘model_onnx.onnx’

engine_path = ‘model_trt.engine’

build_engine(onnx_path1,engine_path,True)

build_engine函数共有三个参数：

onnx_file_path：onnx模型的路径。

engine_file_path：TensorRT模型的路径。

half：是否使用单精度。

单精度的模型速度更快，所以我选择使用单精度。

通过上面的代码就可以完成模型的转化，下面开始实现推理部分，推理分为动态推理和静态推理。

动态推理

---

新建test_trt,py文件，插入代码：

import tensorrt as trt

import pycuda.driver as cuda

import pycuda.autoinit

import numpy as np

import torchvision.transforms as transforms

from PIL import Image

导入需要的包。

def load_engine(engine_path):

TRT_LOGGER = trt.Logger(trt.Logger.WARNING) # INFO

TRT_LOGGER = trt.Logger(trt.Logger.ERROR)

with open(engine_path, ‘rb’) as f, trt.Runtime(TRT_LOGGER) as runtime:

return runtime.deserialize_cuda_engine(f.read())

2. 读取数据，数据处理为可以和网络结构输入对应起来的的shape，数据可增加预处理

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]),

])

定义load_engine函数和get_test_transform函数。

load_engine用于加载TensorRT模型。

get_test_transform实现图像的resize和归一化。

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

image = get_test_transform()(image)

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

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

image = np.array(image)

图片的预处理，和onnx一样，最后转为数组。

path = ‘model_trt.engine’

1. 建立模型，构建上下文管理器

engine = load_engine(path)

context = engine.create_execution_context()

context.active_optimization_profile = 0

3.分配内存空间，并进行数据cpu到gpu的拷贝

动态尺寸，每次都要set一下模型输入的shape，0代表的就是输入，输出根据具体的网络结构而定，可以是0,1,2,3…其中的某个头。

context.set_binding_shape(0, image.shape)

d_input = cuda.mem_alloc(image.nbytes) # 分配输入的内存。

output_shape = context.get_binding_shape(1)

buffer = np.empty(output_shape, dtype=np.float32)

d_output = cuda.mem_alloc(buffer.nbytes) # 分配输出内存。

cuda.memcpy_htod(d_input, image)

bindings = [d_input, d_output]

4.进行推理，并将结果从gpu拷贝到cpu。

context.execute_v2(bindings) # 可异步和同步

cuda.memcpy_dtoh(buffer, d_output)

output = buffer.reshape(output_shape)

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

print(y_pred_binary[0])

输出结果：

静态推理

---

静态推理和动态推理的代码差不多，唯一不同的是不需要

import time

import tensorrt as trt

import pycuda.driver as cuda

import pycuda.autoinit

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

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

二、学习软件

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

三、入门学习视频

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

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

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

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