背景

記錄下onnx轉成TensorRT加速的三種方式

1. 直接使用onnxruntime

在onnxruntime的session初始化的時候第一個provider加入TensorrtExecutionProvider，軟件會自動查找是否支持TensorRT，如果可以就會進行轉換并運行，如果不可以會接著找下一個，也有可能TensorRT跑一半報錯，這就得看環(huán)境什么問題了。

但是這個方法有個很嚴重的問題，它每次運行都要轉換一次，暫時沒看到怎么保存這個轉換完的引擎，對于工程應用初始化太占時間。

2. 使用trtexec.exe

相信想用TensorRT的都已經下載了TensorRT的文件夾了吧，沒下的在這里。在bin文件夾里面有個trtexec.exe。我們可以直接在命令行執(zhí)行這個程序來轉換onnx。
在這邊給個onnx轉tensorrt的示例：

trtexec.exe --onnx=model.onnx --saveEngine=model.trt --fp16

表示轉換model.onnx，保存最終引擎為model.trt（后綴隨意），并使用fp16精度（看個人需求，精度略降，速度提高。并且有些模型使用fp16會出錯）。具體還有一些別的參數(shù)，可以看這個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 可能要根據實際輸出個數(shù)調整下，我因為就一個輸出就直接寫了。我之前碰到過自己不斷申請Stream導致顯存溢出的問題，不知道這個改完會不會還有。

2. 調用session

這一部分是從我自己的代碼里抽象出來的，可能有些問題，如果發(fā)現(xiàn)問題可以評論告訴我

# 初始化
session = TensorRTSession(model_path)
session.get_inputs()
session.get_outputs()
# 推理
session((tensor1, tensor2))

3. 可能遇到的問題

感覺最可能遇到的就一個問題，TensorRT可能會報版本不兼容，我也瞅了好久好久。
選擇自己對于cuda版本的TensorRT之后，開跑發(fā)現(xiàn)老報版本不兼容，但是自己cuda版本和TensorRT看了說明是兼容的，就不是很明白。之前使用C++推理沒這個問題，一開始懷疑是Python版本的TensorRT自己有這個問題。后來發(fā)現(xiàn)是環(huán)境里裝的pytorch自己帶了個CuDNN版本，TensorRT和這個CuDNN版本沖突了。我當時解決方案就是把pytorch改成了cpu版本。只是用得到推理的話可以新建個環(huán)境只裝TensorRT。文章來源地址http://www.zghlxwxcb.cn/news/detail-474121.html

到了這里，關于onnx轉TensorRT使用的三種方式（最終在Python運行）的文章就介紹完了。如果您還想了解更多內容，請在右上角搜索TOY模板網以前的文章或繼續(xù)瀏覽下面的相關文章，希望大家以后多多支持TOY模板網！