代码收藏家 onnx转TensorRT使用的三种方式（最终在Python运行）
背景
记录下onnx转成TensorRT加速的三种方式
1. 直接使用onnxruntime
在onnxruntime的session初始化的时候第一个provider加入TensorrtExecutionProvider，软件会自动查找是否支持TensorRT，如果可以就会进行转换并运行，如果不可以会接着找下一个，也有可能TensorRT跑一半报错，这就得看环境什么问题了。
但是这个方法有个很严重的问题，它每次运行都要转换一次，暂时没看到怎么保存这个转换完的引擎，对于工程应用初始化太占时间。
2. 使用trtexec.exe
相信想用TensorRT的都已经下载了TensorRT的文件夹了吧，没下的在这里。在bin文件夹里面有个trtexec.exe。我们可以直接在命令行执行这个程序来转换onnx。
在这边给个onnx转tensorrt的示例：
trtexec.exe --onnx=model.onnx --saveEngine=model.trt --fp16
表示转换model.onnx，保存最终引擎为model.trt（后缀随意），并使用fp16精度（看个人需求，精度略降，速度提高。并且有些模型使用fp16会出错）。具体还有一些别的参数，可以看这个trtexec.exe的help自己决定。
3. 使用python程序进行转换
参考我之前写的一篇即可
如果需要转换为fp16精度可以加一句话，就这一句什么什么FP16的：
python中使用引擎进行推理
1. 建立trt_session.py
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
# Simple helper data class that's a little nicer to use than a 2-tuple.
class HostDeviceMem(object):
def __init__(self, host_mem, device_mem):
self.host = host_mem
self.device = device_mem
def __str__(self):
return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
def __repr__(self):
return self.__str__()
# Allocates all buffers required for an engine, i.e. host/device inputs/outputs.
def allocate_buffers(engine, context):
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
for i, binding in enumerate(engine):
size = trt.volume(context.get_binding_shape(i))
dtype = trt.nptype(engine.get_binding_dtype(binding))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
# Append the device buffer to device bindings.
bindings.append(int(device_mem))
# Append to the appropriate list.
if engine.binding_is_input(binding):
inputs.append(HostDeviceMem(host_mem, device_mem))
else:
outputs.append(HostDeviceMem(host_mem, device_mem))
return inputs, outputs, bindings, stream
# This function is generalized for multiple inputs/outputs.
# inputs and outputs are expected to be lists of HostDeviceMem objects.
def do_inference(context, bindings, inputs, outputs, stream, batch_size):
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
# Return only the host outputs.
return [out.host for out in outputs]
class TensorRTSession():
def __init__(self, model_path):
f = open(model_path, 'rb')
runtime = trt.Runtime(TRT_LOGGER)
trt.init_libnvinfer_plugins(TRT_LOGGER, '')
self.engine = runtime.deserialize_cuda_engine(f.read())
self.context = self.engine.create_execution_context()
self.inputs_info = None
self.outputs_info = None
self.inputs, self.outputs, self.bindings, self.stream = allocate_buffers(self.engine, self.context)
class nodes_info():
def __init__(self, name, shape):
self.name = name
self.shape = shape
def __call__(self, inputs):
return self.update(inputs)
def get_inputs(self):
inputs_info = []
for i, binding in enumerate(self.engine):
if self.engine.binding_is_input(binding):
shape = self.context.get_binding_shape(i)
inputs_info.append(self.nodes_info(binding, shape))
# print(binding, shape)
self.inputs_info = inputs_info
return inputs_info
def get_outputs(self):
outputs_info = []
for i, binding in enumerate(self.engine):
if not self.engine.binding_is_input(binding):
shape = self.context.get_binding_shape(i)
outputs_info.append(self.nodes_info(binding, shape))
self.outputs_info = outputs_info
return outputs_info
def update(self, input_arr, cuda_ctx=pycuda.autoinit.context):
cuda_ctx.push()
# Do inference
for i in range(len(input_arr)):
self.inputs[i].host = np.ascontiguousarray(input_arr[i])
trt_outputs = do_inference(self.context, bindings=self.bindings, inputs=self.inputs, outputs=self.outputs,
stream=self.stream, batch_size=1)
if cuda_ctx:
cuda_ctx.pop()
trt_outputs = trt_outputs[0].reshape(self.outputs_info[0].shape)
return trt_outputs
最后的trt_outputs 可能要根据实际输出个数调整下，我因为就一个输出就直接写了。我之前碰到过自己不断申请Stream导致显存溢出的问题，不知道这个改完会不会还有。
2. 调用session
这一部分是从我自己的代码里抽象出来的，可能有些问题，如果发现问题可以评论告诉我
# 初始化
session = TensorRTSession(model_path)
session.get_inputs()
session.get_outputs()
# 推理
session((tensor1, tensor2))
3. 可能遇到的问题
感觉最可能遇到的就一个问题，TensorRT可能会报版本不兼容，我也瞅了好久好久。
选择自己对于cuda版本的TensorRT之后，开跑发现老报版本不兼容，但是自己cuda版本和TensorRT看了说明是兼容的，就不是很明白。之前使用C++推理没这个问题，一开始怀疑是Python版本的TensorRT自己有这个问题。后来发现是环境里装的pytorch自己带了个CuDNN版本，TensorRT和这个CuDNN版本冲突了。我当时解决方案就是把pytorch改成了cpu版本。只是用得到推理的话可以新建个环境只装TensorRT。
