TensorRT实战：从模型转换到部署推理的完整指南

1. TensorRT入门为什么选择它如果你正在寻找一种能够让你的深度学习模型在生产环境中飞起来的方法TensorRT绝对是你的不二之选。简单来说TensorRT是NVIDIA推出的高性能推理优化器和运行时引擎专门为NVIDIA GPU设计。它能将你的模型压缩、优化最终在GPU上以惊人的速度运行。我第一次接触TensorRT是在一个实时视频分析项目中。原本用PyTorch直接推理的模型只能跑到30FPS经过TensorRT优化后直接飙到了120FPS效果立竿见影。这种性能提升对于需要低延迟的应用场景比如自动驾驶、工业质检简直是救命稻草。TensorRT主要通过三种方式优化模型层融合将多个操作合并为一个更高效的操作精度校准自动将FP32模型转换为FP16或INT8几乎不损失精度但大幅提升速度内核自动调优为你的特定GPU选择最优的计算内核2. 环境准备安装与验证2.1 下载正确的TensorRT版本首先你需要从NVIDIA官网下载TensorRT。这里有个小坑要注意TensorRT有EAEarly Access和GAGeneral Availability两个版本。作为过来人我强烈建议选择GA版本除非你有特殊需求必须使用最新功能。下载地址https://developer.nvidia.com/tensorrt选择版本时要特别注意CUDA版本兼容性。比如TensorRT-10.0.1.6需要CUDA 12.4如果你用的是CUDA 11.x就得找对应的TensorRT版本。2.2 安装步骤详解下载完成后按照以下步骤安装# 解压文件 tar -xzvf TensorRT-10.0.1.6.Linux.x86_64-gnu.cuda-12.4.tar.gz # 添加库路径到环境变量 echo export LD_LIBRARY_PATH$LD_LIBRARY_PATH:/path/to/TensorRT-10.0.1.6/lib ~/.bashrc source ~/.bashrc # 或者直接复制到CUDA库目录二选一 sudo cp -r ./lib/* /usr/local/cuda/lib64/ # 安装Python包 cd TensorRT-10.0.1.6/python pip install tensorrt-xxx-none-linux_x86_64.whl2.3 验证安装是否成功安装完成后运行以下Python代码验证import tensorrt print(tensorrt.__version__) # 应该输出类似10.0.1.6的版本号 assert tensorrt.Builder(tensorrt.Logger()) # 检查是否能创建Builder如果遇到ImportError: libnvinfer.so.10: cannot open shared object file这样的错误说明环境变量没设置好。回头检查LD_LIBRARY_PATH是否包含了TensorRT的lib目录。3. 模型转换从PyTorch到TensorRT3.1 方法一PyTorch → ONNX → TensorRT这是最常用的转换路径兼容性最好。让我们以ResNet18为例import torch import torchvision # 加载模型 model torchvision.models.resnet18(pretrainedFalse).eval().cuda() # 创建虚拟输入 dummy_input torch.randn(1, 3, 224, 224, devicecuda) # 导出ONNX模型 torch.onnx.export( model, dummy_input, resnet18.onnx, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size}, output: {0: batch_size} }, opset_version14 ) # 验证ONNX模型 import onnx onnx_model onnx.load(resnet18.onnx) onnx.checker.check_model(onnx_model)接下来使用TensorRT的trtexec工具转换ONNX到TensorRT引擎/path/to/TensorRT/bin/trtexec \ --onnxresnet18.onnx \ --saveEngineresnet18.trt \ --fp16 \ --shapesinput:1x3x224x2243.2 方法二直接使用torch2trt如果你想要更简单的转换方式可以试试torch2trtfrom torch2trt import torch2trt # 转换模型 model_trt torch2trt(model, [dummy_input]) # 保存模型 torch.save(model_trt.state_dict(), resnet18_trt.pth) # 加载模型 from torch2trt import TRTModule model_trt TRTModule() model_trt.load_state_dict(torch.load(resnet18_trt.pth))这种方法简单直接但灵活性不如ONNX路径某些自定义操作可能不支持。4. 性能优化技巧4.1 精度选择FP32/FP16/INT8TensorRT支持三种精度模式FP32最高精度速度最慢FP16精度损失很小速度提升明显推荐首选INT8需要校准精度可能下降但速度最快在trtexec中通过--fp16或--int8启用trtexec --onnxmodel.onnx --saveEnginemodel_fp16.trt --fp16 trtexec --onnxmodel.onnx --saveEnginemodel_int8.trt --int8对于INT8你还需要提供校准数据# 创建校准器 calibrator trt.Int8EntropyCalibrator2( calibration_data, batch_size32) # 构建时指定校准器 builder.int8_calibrator calibrator4.2 动态形状支持如果你的模型需要处理不同大小的输入可以启用动态形状profile builder.create_optimization_profile() profile.set_shape( input, # 输入名称 (1, 3, 224, 224), # 最小形状 (8, 3, 224, 224), # 最优形状 (32, 3, 224, 224) # 最大形状 ) config.add_optimization_profile(profile)4.3 层融合与图优化TensorRT会自动进行图优化但你也可以通过配置控制config builder.create_builder_config() config.set_flag(trt.BuilderFlag.FP16) # 启用FP16 config.set_flag(trt.BuilderFlag.STRICT_TYPES) # 强制使用指定精度 config.max_workspace_size 1 30 # 设置工作空间大小5. 部署实战Python与C示例5.1 Python部署代码import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit # 加载引擎 TRT_LOGGER trt.Logger(trt.Logger.INFO) with open(resnet18.trt, rb) as f, \ trt.Runtime(TRT_LOGGER) as runtime, \ runtime.deserialize_cuda_engine(f.read()) as engine, \ engine.create_execution_context() as context: # 准备输入 input_data np.random.randn(1, 3, 224, 224).astype(np.float32) # 分配设备内存 d_input cuda.mem_alloc(input_data.nbytes) d_output cuda.mem_alloc(input_data.nbytes) # 创建流 stream cuda.Stream() # 拷贝数据到设备 cuda.memcpy_htod_async(d_input, input_data, stream) # 执行推理 context.execute_async_v3(bindings[int(d_input), int(d_output)], stream_handlestream.handle) # 拷贝结果回主机 output_data np.empty_like(input_data) cuda.memcpy_dtoh_async(output_data, d_output, stream) # 同步 stream.synchronize() print(推理结果:, output_data)5.2 C部署代码#include NvInfer.h #include NvOnnxParser.h #include cuda_runtime_api.h // 加载引擎 std::ifstream engineFile(resnet18.trt, std::ios::binary); engineFile.seekg(0, std::ios::end); size_t size engineFile.tellg(); engineFile.seekg(0, std::ios::beg); std::vectorchar engineData(size); engineFile.read(engineData.data(), size); nvinfer1::IRuntime* runtime nvinfer1::createInferRuntime(logger); nvinfer1::ICudaEngine* engine runtime-deserializeCudaEngine(engineData.data(), size); nvinfer1::IExecutionContext* context engine-createExecutionContext(); // 准备输入输出缓冲区 void* buffers[2]; cudaMalloc(buffers[0], inputSize * sizeof(float)); cudaMalloc(buffers[1], outputSize * sizeof(float)); // 创建CUDA流 cudaStream_t stream; cudaStreamCreate(stream); // 拷贝输入数据 cudaMemcpyAsync(buffers[0], inputData, inputSize * sizeof(float), cudaMemcpyHostToDevice, stream); // 执行推理 context-enqueueV2(buffers, stream, nullptr); // 拷贝输出数据 cudaMemcpyAsync(outputData, buffers[1], outputSize * sizeof(float), cudaMemcpyDeviceToHost, stream); // 同步 cudaStreamSynchronize(stream);6. 常见问题与解决方案6.1 模型转换失败问题转换时出现Unsupported ONNX opset version错误。解决尝试降低opset版本。PyTorch导出ONNX时设置opset_version11或更低torch.onnx.export(..., opset_version11)6.2 推理结果不正确问题TensorRT推理结果与原始框架不一致。解决检查输入数据预处理是否一致确保所有操作都被支持使用polygraphy工具检查尝试禁用优化builder_config.set_flag(trt.BuilderFlag.DISABLE_TIMING_CACHE)6.3 性能不如预期问题优化后速度提升不明显。解决检查是否真正使用了FP16/INT8context.get_tensor_type(tensor_name)增加工作空间大小config.max_workspace_size 1 30尝试不同的优化配置config.set_flag(trt.BuilderFlag.TF32)在实际项目中我发现TensorRT的性能调优是个迭代过程。建议先用小批量数据测试各种配置找到最优组合后再进行大规模部署。记得每次修改配置后都要重新构建引擎因为TensorRT的优化是基于当前配置的。