时间:2023/11/01
说明:本指南针对在装有NVIDIA显卡的Windows10系统的计算机上,安装TensorRT推理加速工具,将pytorch中导出的onnx模型转换为trt模型文件,并在python语言中进行模型推理。
TensorRT官方安装指南:https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html#installing-zip
注意:
将ONNX模型转换成TensorRT的最简单的方法就是使用{TensorRT}/bin
下的命令行工具trtexec
:
trtexec --onnx=model.onnx --saveEngine=model.trt
这里的--onnx
和--saveEngine
分别代表onnx模型的路径和保存trt模型的路径。此外,再介绍两个比较常用的trtexec
命令行工具参数:
--explicitBatch
:告诉trtexec在优化时固定输入的 batch size(将从onnx文件中推断batch size的具体值,即与导出onnx文件时传入的batch size一致)。当确定模型的输入batch size时,推荐采用此参数,因为固定batch size大小可以使得trtexec进行额外的优化,且省去了指定“优化配置文件”这一额外步骤(采用动态batch size时需要提供“优化配置文件”来指定希望接收的可能的batch size大小的范围);--fp16
:采用FP16精度,通过牺牲部分模型准确率来简化模型(减少显存占用和加速模型推理)。TensorRT支持TF32/FP32/FP16/INT8多种精度(具体还要看GPU是否支持)。FP32是多数框架训练模型的默认精度,FP16对模型推理速度和显存占用有较大优化,且准确率损失往往可以忽略不计。INT8进一步牺牲了准确率,同时也进一步降低模型的延迟和显存要求,但需要额外的步骤来仔细校准,来使其精度损耗较小。使用Python API运行TensorRT模型推理需要安装pycuda
包:
pip install pycuda
注:若pycuda安装失败,尝试到https://www.lfd.uci.edu/~gohlke/pythonlibs/#pycuda 下载python版本对应的最新的本地安装文件安装
然后参照官方给的示例代码运行TensorRT模型推理:tutorial-runtime.ipynb
下面给出Unet语义分割模型运行tensorrt推理的主要代码:
(1) 导入TensorRT推理需要的库
import pycuda.driver as cuda
import pycuda.autoinit
import tensorrt as trt
(2) 加载TensorRT引擎
def load_engine(self, engine_file_path):
TRT_LOGGER = trt.Logger()
assert os.path.exists(engine_file_path)
print("Reading engine from file {}".format(engine_file_path))
with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
return runtime.deserialize_cuda_engine(f.read())
(3) 进行推理
def engine_infer(self, engine, input_image):
"""
engine: load_engine函数返回的trt模型引擎
input_image: 模型推理输入图像,尺寸为(batch_size, channel, height, width)
output:Unet模型推理的结果,尺寸为(batch_size, class_num, height, width)
"""
batch_size = input_image.shape[0]
image_channel = input_image.shape[1]
image_height = input_image.shape[2]
image_width = input_image.shape[3]
with engine.create_execution_context() as context:
# Set input shape based on image dimensions for inference
context.set_binding_shape(engine.get_binding_index("input"), (batch_size, image_channel, image_height, image_width))
# Allocate host and device buffers
bindings = []
for binding in engine:
binding_idx = engine.get_binding_index(binding)
size = trt.volume(context.get_binding_shape(binding_idx))
dtype = trt.nptype(engine.get_binding_dtype(binding))
if engine.binding_is_input(binding):
input_buffer = np.ascontiguousarray(input_image)
input_memory = cuda.mem_alloc(input_image.nbytes)
bindings.append(int(input_memory))
else:
output_buffer = cuda.pagelocked_empty(size, dtype)
output_memory = cuda.mem_alloc(output_buffer.nbytes)
bindings.append(int(output_memory))
stream = cuda.Stream()
# Transfer input data to the GPU.
cuda.memcpy_htod_async(input_memory, input_buffer, stream)
# Run inference
context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
# Transfer prediction output from the GPU.
cuda.memcpy_dtoh_async(output_buffer, output_memory, stream)
# Synchronize the stream
stream.synchronize()
output = np.reshape(output_buffer, (batch_size, self.num_classes, image_height, image_width))
return output
不同显卡(不同GPU),其核心数量、频率、架构、设计都是不一样的,TensorRT需要对特定的硬件进行优化,不同硬件之间的优化是不能共享的。
支持FP32、FP16、INT8、TF32等,这几种类型都比较常用。