当前位置：首页 > news >正文

TensorRT安装部署指南（Windows10）

news 来源：原创 2024/5/20 17:20:58

TensorRT安装部署指南

时间：2023/11/01
说明：本指南针对在装有NVIDIA显卡的Windows10系统的计算机上，安装TensorRT推理加速工具，将pytorch中导出的onnx模型转换为trt模型文件，并在python语言中进行模型推理。

一、TensorRT安装

TensorRT官方安装指南：https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html#installing-zip

注意：

官方给出了pip、RPM、Tar文件、Zip文件等多种安装方式。需要注意，通过zip文件安装是目前 Windows 的唯一选项。Zip文件从官网下载。下载完成后，只需要按照官方安装指南中给出的步骤进行安装和环境配置。
安装TensorRT前需要确保已经安装对应的CUDA和cuDNN依赖。CUDA和cuDNN的安装方法可以参考：CUDA和cudnn安装教程。

二、将ONNX模型转换为TensorRT引擎

将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模型进行推理

使用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 inferencecontext.set_binding_shape(engine.get_binding_index("input"), (batch_size, image_channel, image_height, image_width))# Allocate host and device buffersbindings = []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 inferencecontext.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 streamstream.synchronize()output = np.reshape(output_buffer, (batch_size, self.num_classes, image_height, image_width))return output

注意事项

（1）TensorRT是硬件相关的

不同显卡(不同GPU)，其核心数量、频率、架构、设计都是不一样的，TensorRT需要对特定的硬件进行优化，不同硬件之间的优化是不能共享的。

（2）TensorRT支持哪几种权重精度

支持FP32、FP16、INT8、TF32等，这几种类型都比较常用。

FP32：单精度浮点型，深度学习中最常见的数据格式，训练推理都会用到；
FP16：半精度浮点型，相比FP32占用内存减少一半，有相应的指令值，速度比FP32要快很多；
TF32：第三代Tensor Core支持的一种数据类型，是一种截短的 Float32 数据格式，将FP32中23个尾数位截短为10bits，而指数位仍为8bits，总长度为19(=1+8 +10)。保持了与FP16同样的精度(尾数位都是 10 位），同时还保持了FP32的动态范围指数位都是8位)；
INT8：整型，相比FP16占用内存减小一半，有相应的指令集，模型量化后可以利用INT8进行加速。