1. 从零开始TensorRT Python API 能为你做什么如果你正在为深度学习模型推理速度慢而头疼或者想把训练好的模型快速部署到生产环境那TensorRT的Python API绝对是你应该立刻拿起来的工具。我用了这么多年最大的感受就是它能把你的模型推理速度提升好几倍而且用起来比想象中简单得多。简单来说TensorRT是英伟达专门为自家GPU打造的推理优化引擎。你可以把它想象成一个“模型编译器”——它会把你的PyTorch、TensorFlow等框架训练出来的模型重新编译成专门针对你当前显卡硬件优化的执行计划。这个优化过程包括了算子融合、精度校准、内存优化等一系列“黑科技”。我实测过不少模型像ResNet、YOLO这类常见网络优化后速度提升3-5倍是家常便饭有些模型甚至能快上10倍。那么Python API在这个生态里扮演什么角色呢它让你能用最熟悉的Python语言完成从模型导出、优化到部署的整个流程。你不用去碰复杂的C代码也不用折腾繁琐的编译环境pip安装一下几行代码就能跑起来。这对于算法工程师和快速原型开发来说简直是福音。我见过很多团队一开始用C API折腾得死去活来后来切换到Python API开发效率直接翻倍。不过我得提醒你TensorRT的加速效果和你的模型结构、硬件配置强相关。如果你的模型里有很多自定义算子或者动态控制流优化起来可能会麻烦一些。但好消息是现在TensorRT对主流模型的支持已经非常好了ONNX生态也越来越成熟大部分情况下你都能获得不错的加速比。2. 环境准备别在第一步就踩坑开始之前咱们先把环境搭好。这一步看似简单但很多人都在这里栽了跟头。我根据自己踩过的坑给你整理了一份最稳妥的安装方案。首先TensorRT对CUDA和cuDNN的版本有严格要求。我建议你直接用英伟达官方提供的Docker镜像这是最省心的方式。但如果你需要在本地环境安装那就得仔细核对版本兼容性了。以最新的TensorRT 10.x为例它通常需要CUDA 12.x和cuDNN 8.9.x。你可以用下面的命令检查现有环境nvcc --version # 查看CUDA版本 cat /usr/local/cuda/version.txt # 另一种查看CUDA版本的方式如果版本不匹配建议先升级或降级CUDA。别想着凑合用版本不兼容会导致各种诡异的问题我在这上面浪费过不少时间。接下来安装TensorRT Python包。现在最简单的方法就是直接用pip安装pip install tensorrt这个tensorrt包其实是个“元包”它会自动帮你安装tensorrt-libs、tensorrt-bindings等依赖。但这里有个坑pip安装的版本可能不是最新的而且有时候会因为网络问题安装失败。如果遇到问题你可以尝试从英伟达官网下载对应版本的tar包手动安装Python wheel文件。安装完成后验证一下是否成功import tensorrt as trt print(trt.__version__) # 应该输出类似10.15.1这样的版本号 # 创建一个简单的logger测试 logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) print(TensorRT Python API导入成功)如果这一步没问题恭喜你环境搭建完成了。但我还得提醒几个常见问题第一确保你的Python环境是64位的32位Python不支持第二如果你在Windows上可能需要手动设置一些环境变量第三如果你要用ONNX解析器还需要额外安装onnx和onnx-graphsurgeon这两个包。3. 模型导出把PyTorch模型变成ONNX格式现在环境准备好了咱们开始第一步把训练好的模型导出为ONNX格式。ONNX就像模型界的“普通话”几乎所有框架都能听懂它。TensorRT通过ONNX解析器来理解你的模型结构所以导出ONNX是必经之路。我以PyTorch为例给你展示一个完整的导出流程。假设我们有一个简单的图像分类模型import torch import torch.nn as nn import torch.onnx class SimpleCNN(nn.Module): def __init__(self, num_classes10): super(SimpleCNN, self).__init__() self.conv1 nn.Conv2d(3, 32, kernel_size3, padding1) self.relu1 nn.ReLU() self.pool1 nn.MaxPool2d(2) self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) self.relu2 nn.ReLU() self.pool2 nn.MaxPool2d(2) self.fc nn.Linear(64 * 8 * 8, num_classes) # 假设输入是32x32 def forward(self, x): x self.pool1(self.relu1(self.conv1(x))) x self.pool2(self.relu2(self.conv2(x))) x x.view(x.size(0), -1) x self.fc(x) return x # 创建模型实例并加载权重 model SimpleCNN() model.load_state_dict(torch.load(model_weights.pth)) model.eval() # 重要切换到推理模式 # 准备一个示例输入 dummy_input torch.randn(1, 3, 32, 32) # batch_size1, channels3, height32, width32 # 导出ONNX input_names [input] output_names [output] onnx_path simple_cnn.onnx torch.onnx.export( model, dummy_input, onnx_path, export_paramsTrue, # 导出权重参数 opset_version13, # ONNX算子集版本建议用11或更高 do_constant_foldingTrue, # 常量折叠优化 input_namesinput_names, output_namesoutput_names, dynamic_axes{ input: {0: batch_size}, # 支持动态batch output: {0: batch_size} } ) print(f模型已导出到 {onnx_path})这里有几个关键点需要特别注意。首先是model.eval()这行代码千万不能忘。PyTorch的某些层比如Dropout、BatchNorm在训练和推理时的行为不一样如果不切换到eval模式导出的ONNX模型可能包含训练专用的算子。其次是opset_version我建议用11或更高版本。太老的版本可能不支持一些新算子太新的版本又可能和TensorRT的解析器不兼容。目前TensorRT 10.x对ONNX opset 13支持得很好。dynamic_axes参数是可选的但如果你想让模型支持动态batch或动态尺寸一定要在这里指定。比如上面代码中我把batch维度设成了动态的这样同一个引擎就能处理不同batch size的输入了。这在部署时非常有用特别是当你的服务需要处理不同并发请求时。导出完成后我强烈建议你用Netron一个开源的可视化工具打开ONNX文件检查一下。看看输入输出维度对不对算子有没有被正确识别。我遇到过不少情况PyTorch导出时一切正常但TensorRT解析时却报错往往就是因为ONNX文件里有些奇怪的节点。4. 构建引擎把ONNX变成TensorRT的“执行计划”拿到ONNX文件后下一步就是构建TensorRT引擎。这个过程有点像编译C代码ONNX是源代码TensorRT是编译器最后生成的是针对特定硬件优化的二进制执行计划。让我给你展示一个完整的构建脚本我会边写代码边解释每个参数的作用import tensorrt as trt import os def build_engine_from_onnx(onnx_path, engine_path, max_batch_size1, fp16_modeTrue, max_workspace_size1 30): # 1GB 从ONNX文件构建TensorRT引擎 参数 - onnx_path: ONNX模型文件路径 - engine_path: 要保存的引擎文件路径 - max_batch_size: 最大batch size影响内存分配 - fp16_mode: 是否启用FP16精度能显著减少显存和加速 - max_workspace_size: 构建引擎时的临时工作空间大小 # 1. 创建logger用于记录构建过程中的信息 # TRT的日志级别有VERBOSE, INFO, WARNING, ERROR, INTERNAL_ERROR # 调试时可以用INFO生产环境建议用WARNING或ERROR logger trt.Logger(trt.Logger.WARNING) # 2. 创建构建器 builder trt.Builder(logger) # 3. 创建网络定义 # EXPLICIT_BATCH标志表示我们使用显式batch维度 # 这是ONNX解析器要求的 network_flags 1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) network builder.create_network(network_flags) # 4. 创建ONNX解析器 parser trt.OnnxParser(network, logger) # 5. 解析ONNX文件 with open(onnx_path, rb) as f: onnx_data f.read() if not parser.parse(onnx_data): print(ONNX解析失败错误信息) for i in range(parser.num_errors): print(f 错误 {i}: {parser.get_error(i)}) return None print(ONNX解析成功) print(f网络层数{network.num_layers}) print(f输入数量{network.num_inputs}) print(f输出数量{network.num_outputs}) # 6. 创建构建配置 config builder.create_builder_config() # 设置工作空间大小 # 这个空间用于层实现时的临时存储 # 太小可能导致某些层无法找到实现太大会浪费内存 config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, max_workspace_size) # 7. 设置精度模式 if fp16_mode and builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) print(启用FP16精度模式) else: print(使用FP32精度模式) # 8. 设置动态形状如果需要 # 如果你的模型需要支持动态输入尺寸需要在这里配置 profile builder.create_optimization_profile() # 获取输入名称 input_name network.get_input(0).name # 设置最小、最优、最大形状 # 这里假设输入是 [batch, channel, height, width] min_shape (1, 3, 32, 32) # 最小输入尺寸 opt_shape (max_batch_size, 3, 32, 32) # 最优尺寸推理时最常用 max_shape (max_batch_size, 3, 224, 224) # 最大支持尺寸 profile.set_shape(input_name, min_shape, opt_shape, max_shape) config.add_optimization_profile(profile) # 9. 构建引擎 print(开始构建引擎...这可能需要几分钟) serialized_engine builder.build_serialized_network(network, config) if serialized_engine is None: print(引擎构建失败) return None # 10. 保存引擎到文件 with open(engine_path, wb) as f: f.write(serialized_engine) print(f引擎构建成功已保存到 {engine_path}) # 打印一些统计信息 print(f引擎序列化大小{len(serialized_engine) / 1024 / 1024:.2f} MB) return serialized_engine # 使用示例 if __name__ __main__: onnx_file simple_cnn.onnx engine_file simple_cnn.engine # 构建引擎 engine build_engine_from_onnx( onnx_file, engine_file, max_batch_size16, # 支持最大batch为16 fp16_modeTrue, # 启用FP16加速 max_workspace_size2 30 # 2GB工作空间 ) if engine: print(引擎构建完成可以用于推理了)这段代码有几个关键点需要特别关注。首先是工作空间大小max_workspace_size这个参数决定了TensorRT在优化过程中能使用多少临时显存。如果设得太小TensorRT可能无法为某些层找到最优的实现如果设得太大又会浪费显存。我一般从1GB开始试如果构建失败就适当调大。FP16模式是个好东西它能将模型精度从32位浮点数降到16位通常能带来1.5-2倍的加速同时显存占用减半。但要注意不是所有显卡都支持FP16加速你可以通过builder.platform_has_fast_fp16来检查。另外FP16可能会带来轻微精度损失对于精度敏感的任务比如医疗影像要谨慎使用。动态形状配置是实际部署中的关键。想象一下你的模型可能需要处理不同尺寸的图片或者batch size需要根据请求量动态调整。通过设置最小、最优、最大三个形状TensorRT会为这个范围内的所有尺寸生成优化策略。注意范围设得越大构建时间越长引擎文件也越大。构建过程可能会比较耗时特别是对于大模型。在我的RTX 4090上构建一个ResNet-50大概需要2-3分钟更大的模型可能需要10分钟以上。这时候你可以去喝杯咖啡或者看看构建日志里有没有警告信息。5. 推理部署让优化后的引擎跑起来引擎构建好了现在到了最激动人心的环节实际推理。这部分代码虽然不长但每个细节都关系到推理的效率和稳定性。我结合自己的实战经验给你写了一个既健壮又高效的推理类import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import numpy as np from typing import Dict, List, Optional import time class TensorRTInference: def __init__(self, engine_path: str, gpu_id: int 0): 初始化TensorRT推理引擎 参数 - engine_path: 引擎文件路径 - gpu_id: 使用的GPU设备ID # 设置当前GPU cuda.init() self.device cuda.Device(gpu_id) self.context self.device.make_context() # 创建logger self.logger trt.Logger(trt.Logger.WARNING) # 加载并反序列化引擎 with open(engine_path, rb) as f: engine_data f.read() self.runtime trt.Runtime(self.logger) self.engine self.runtime.deserialize_cuda_engine(engine_data) # 创建执行上下文 self.context self.engine.create_execution_context() # 分配输入输出内存 self._setup_bindings() # 创建CUDA流用于异步执行 self.stream cuda.Stream() print(f引擎加载成功) print(f 输入数量{len(self.inputs)}) print(f 输出数量{len(self.outputs)}) print(f 最大batch size{self.engine.max_batch_size}) def _setup_bindings(self): 设置输入输出绑定和分配内存 self.inputs [] self.outputs [] self.bindings [] # 遍历所有绑定输入输出 for i in range(self.engine.num_bindings): binding_name self.engine.get_binding_name(i) dtype self.engine.get_binding_dtype(i) shape self.engine.get_binding_shape(i) is_input self.engine.binding_is_input(i) # 计算需要分配的内存大小 volume trt.volume(shape) if self.engine.has_implicit_batch_dimension: volume * self.engine.max_batch_size # 根据数据类型确定numpy dtype和元素大小 if dtype trt.DataType.FLOAT: np_dtype np.float32 element_size 4 elif dtype trt.DataType.HALF: np_dtype np.float16 element_size 2 elif dtype trt.DataType.INT32: np_dtype np.int32 element_size 4 elif dtype trt.DataType.INT8: np_dtype np.int8 element_size 1 else: raise TypeError(f不支持的TensorRT数据类型: {dtype}) total_memory volume * element_size # 分配锁页主机内存 host_mem cuda.pagelocked_empty(volume, dtypenp_dtype) # 分配设备内存 device_mem cuda.mem_alloc(total_memory) # 保存绑定信息 binding_info { name: binding_name, host: host_mem, device: device_mem, shape: shape, dtype: np_dtype, is_input: is_input, volume: volume } if is_input: self.inputs.append(binding_info) else: self.outputs.append(binding_info) self.bindings.append(int(device_mem)) print(f {输入 if is_input else 输出}: {binding_name}) print(f 形状: {shape}, 数据类型: {dtype}, 元素数量: {volume}) def infer(self, input_data: Dict[str, np.ndarray]) - Dict[str, np.ndarray]: 执行推理 参数 - input_data: 字典键为输入名称值为numpy数组 返回 - 字典键为输出名称值为numpy数组 # 1. 准备输入数据 for input_info in self.inputs: name input_info[name] if name not in input_data: raise ValueError(f缺少输入: {name}) data input_data[name] # 检查输入形状是否匹配 expected_shape tuple(input_info[shape]) if data.shape ! expected_shape: # 尝试reshape如果元素总数相同 if data.size input_info[volume]: data data.reshape(expected_shape) else: raise ValueError( f输入 {name} 形状不匹配。期望 {expected_shape}得到 {data.shape} ) # 检查数据类型 if data.dtype ! input_info[dtype]: data data.astype(input_info[dtype]) # 复制数据到锁页内存 np.copyto(input_info[host], data.ravel()) # 异步复制到设备 cuda.memcpy_htod_async( input_info[device], input_info[host], self.stream ) # 2. 设置动态形状如果有 # 如果引擎支持动态形状需要在这里设置实际的输入形状 for i, input_info in enumerate(self.inputs): if -1 in input_info[shape]: # 动态维度 actual_shape input_data[input_info[name]].shape self.context.set_binding_shape(i, actual_shape) # 3. 执行推理 self.context.execute_async_v2( bindingsself.bindings, stream_handleself.stream.handle ) # 4. 将输出从设备复制回主机 results {} for output_info in self.outputs: cuda.memcpy_dtoh_async( output_info[host], output_info[device], self.stream ) # 等待异步操作完成 self.stream.synchronize() # 根据实际形状reshape输出 if self.context.all_binding_shapes_specified: output_shape self.context.get_binding_shape( self.engine.get_binding_index(output_info[name]) ) results[output_info[name]] output_info[host].reshape(output_shape) else: # 使用原始形状 results[output_info[name]] output_info[host].reshape(output_info[shape]) return results def benchmark(self, input_data: Dict[str, np.ndarray], warmup: int 10, iterations: int 100): 性能基准测试 参数 - input_data: 测试输入数据 - warmup: 预热迭代次数 - iterations: 测试迭代次数 print(开始性能测试...) # 预热 print(f预热 {warmup} 次...) for _ in range(warmup): _ self.infer(input_data) # 正式测试 print(f运行 {iterations} 次推理...) latencies [] for i in range(iterations): start_time time.perf_counter() _ self.infer(input_data) end_time time.perf_counter() latencies.append((end_time - start_time) * 1000) # 转换为毫秒 if (i 1) % 20 0: print(f 已完成 {i 1}/{iterations} 次) # 统计结果 avg_latency np.mean(latencies) min_latency np.min(latencies) max_latency np.max(latencies) std_latency np.std(latencies) print(\n性能测试结果) print(f 平均延迟: {avg_latency:.2f} ms) print(f 最小延迟: {min_latency:.2f} ms) print(f 最大延迟: {max_latency:.2f} ms) print(f 标准差: {std_latency:.2f} ms) print(f FPS: {1000 / avg_latency:.2f}) return { avg_latency: avg_latency, min_latency: min_latency, max_latency: max_latency, std_latency: std_latency, fps: 1000 / avg_latency } def __del__(self): 清理资源 if hasattr(self, context): self.context.pop() del self.context if hasattr(self, stream): del self.stream # 使用示例 if __name__ __main__: # 初始化推理引擎 inferencer TensorRTInference(simple_cnn.engine) # 准备输入数据 # 注意形状需要和构建引擎时设置的一致 input_data { input: np.random.randn(1, 3, 32, 32).astype(np.float32) } # 执行推理 print(执行推理...) outputs inferencer.infer(input_data) # 打印输出 for name, tensor in outputs.items(): print(f输出 {name}: 形状{tensor.shape}, 数据类型{tensor.dtype}) print(f 前5个值: {tensor.flatten()[:5]}) # 性能测试 print(\n进行性能基准测试...) stats inferencer.benchmark(input_data, warmup10, iterations100)这个推理类我打磨了很久里面有很多实战中总结的经验。首先是内存管理我使用了锁页内存pinned memory这种内存的特点是CPU和GPU之间的数据传输速度更快。对于推理服务来说这能显著降低数据搬运的开销。异步执行是另一个关键优化。你看execute_async_v2和memcpy_htod_async这些调用它们都是非阻塞的。这意味着数据拷贝和计算可以重叠进行GPU利用率更高。配合CUDA流的同步机制既能保证正确性又能最大化性能。动态形状支持在实际部署中特别有用。比如你的服务可能同时收到不同尺寸的图片或者batch size需要动态调整。通过set_binding_shape方法你可以在运行时指定实际的输入形状TensorRT会自动选择最优的kernel。性能测试部分我建议你每次都跑一下。通过基准测试你不仅能了解模型的推理速度还能发现性能瓶颈在哪里。我遇到过的情况包括输入数据预处理太慢、内存拷贝开销太大、甚至GPU温度过高导致降频。这些都需要具体问题具体分析。6. 高级技巧动态输入与多精度优化掌握了基础流程后咱们来聊聊一些高级技巧。这些技巧能让你的TensorRT部署更加灵活高效特别是在生产环境中。6.1 动态输入处理实战动态输入是实际部署中的常见需求。比如你的图像分类服务用户可能上传各种尺寸的图片。TensorRT通过优化配置文件Optimization Profile来支持这个功能。让我给你展示一个完整的动态输入处理示例def build_dynamic_engine(onnx_path, engine_path): 构建支持动态输入的引擎 logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) # 解析ONNX parser trt.OnnxParser(network, logger) with open(onnx_path, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) # 创建优化配置文件 profile builder.create_optimization_profile() # 假设模型有一个输入叫input input_name network.get_input(0).name # 设置动态范围 # 最小尺寸batch1, 224x224 # 最优尺寸batch8, 224x224推理时最常用的尺寸 # 最大尺寸batch32, 512x512 profile.set_shape( input_name, min(1, 3, 224, 224), opt(8, 3, 224, 224), max(32, 3, 512, 512) ) config.add_optimization_profile(profile) # 如果需要多个输入都支持动态可以创建多个profile # 但通常一个profile就够用了 # 构建引擎 engine builder.build_engine(network, config) # 保存引擎 with open(engine_path, wb) as f: f.write(engine.serialize()) return engine # 使用动态引擎推理 def infer_with_dynamic_shape(inferencer, batch_size, height, width): 使用动态形状进行推理 # 创建随机输入数据 input_data np.random.randn(batch_size, 3, height, width).astype(np.float32) # 设置实际的输入形状 input_name inferencer.inputs[0][name] # 对于动态形状需要在推理前设置binding shape input_index inferencer.engine.get_binding_index(input_name) inferencer.context.set_binding_shape(input_index, input_data.shape) # 重新分配输出内存因为输出形状可能变了 # 这里需要根据新的输出形状重新分配内存 # 实际代码中需要更复杂的内存管理 # 执行推理 outputs inferencer.infer({input_name: input_data}) return outputs动态输入有几个需要注意的地方。首先TensorRT会为每个优化配置文件生成多组kernel所以引擎文件会变大构建时间也会变长。其次你需要在推理时通过set_binding_shape告诉TensorRT实际的输入尺寸否则它会使用默认尺寸。我建议你根据实际业务场景来设置动态范围。范围设得太宽引擎会变得臃肿范围设得太窄又可能无法处理某些输入。通常我会分析历史请求数据找到最常见的尺寸范围。6.2 多精度优化策略精度优化是TensorRT的强项。除了FP32和FP16TensorRT还支持INT8量化能带来更大的加速比。但量化是有代价的你需要仔细权衡精度损失和速度提升。def build_with_mixed_precision(onnx_path, engine_path, precision_modefp16): 构建多精度引擎 logger trt.Logger(trt.Logger.INFO) # 用INFO级别看更多细节 builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser trt.OnnxParser(network, logger) with open(onnx_path, rb) as f: parser.parse(f.read()) config builder.create_builder_config() # 设置工作空间 config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 30) # 根据精度模式设置标志 if precision_mode fp16: if builder.platform_has_fast_fp16: config.set_flag(trt.BuilderFlag.FP16) print(启用FP16精度) else: print(当前平台不支持FP16回退到FP32) elif precision_mode int8: if builder.platform_has_fast_int8: config.set_flag(trt.BuilderFlag.INT8) print(启用INT8精度) # INT8需要校准数据 # 这里使用简单的随机数据作为示例 # 实际应用中应该使用有代表性的数据集 class RandomCalibrator(trt.IInt8EntropyCalibrator2): def __init__(self, batch_size32): super().__init__() self.batch_size batch_size self.shape (batch_size, 3, 224, 224) self.batch_count 100 # 使用100个batch进行校准 self.current_batch 0 def get_batch_size(self): return self.batch_size def get_batch(self, names): if self.current_batch self.batch_count: # 生成随机校准数据 batch np.random.randn(*self.shape).astype(np.float32) # 分配设备内存 device_memory cuda.mem_alloc(batch.nbytes) cuda.memcpy_htod(device_memory, batch) self.current_batch 1 return [int(device_memory)] else: return None def read_calibration_cache(self): # 可以读取之前保存的校准缓存 return None def write_calibration_cache(self, cache): # 保存校准缓存供下次使用 with open(calibration.cache, wb) as f: f.write(cache) calibrator RandomCalibrator() config.int8_calibrator calibrator else: print(当前平台不支持INT8回退到FP16或FP32) # 构建引擎 print(开始构建引擎...) engine builder.build_engine(network, config) if engine: with open(engine_path, wb) as f: f.write(engine.serialize()) print(f引擎构建完成精度模式{precision_mode}) return engine # 精度测试比较 def compare_precision(original_model, trt_fp32, trt_fp16, trt_int8, test_data): 比较不同精度下的推理结果 results {} # 原始模型推理作为基准 with torch.no_grad(): original_output original_model(test_data) results[original] original_output.numpy() # TensorRT FP32推理 fp32_output trt_fp32.infer({input: test_data.numpy()}) results[fp32] fp32_output[output] # TensorRT FP16推理 fp16_output trt_fp16.infer({input: test_data.numpy()}) results[fp16] fp16_output[output] # TensorRT INT8推理 int8_output trt_int8.infer({input: test_data.numpy()}) results[int8] int8_output[output] # 计算与原始结果的差异 print(精度比较结果) for key in [fp32, fp16, int8]: diff np.abs(results[original] - results[key]).mean() print(f {key} 与原始结果的平均绝对误差{diff:.6f}) return resultsINT8量化是个技术活校准数据的选择直接影响量化质量。我建议你使用有代表性的真实数据而不是随机数据。校准数据应该覆盖模型可能遇到的各种输入情况。如果校准数据不够有代表性量化后的模型在某些输入上可能会出现较大误差。另外TensorRT提供了多种校准方法熵校准、最小最大校准等。熵校准通常效果更好但计算量更大。你可以根据实际需求选择。6.3 性能调优实战构建好引擎后还可以通过一些技巧进一步提升性能。这里我分享几个实测有效的优化方法class OptimizedTensorRTInference(TensorRTInference): 优化版的TensorRT推理器 def __init__(self, engine_path, gpu_id0, use_cuda_graphFalse): super().__init__(engine_path, gpu_id) self.use_cuda_graph use_cuda_graph self.cuda_graph None self.graph_executable None if use_cuda_graph: self._capture_cuda_graph() def _capture_cuda_graph(self): 捕获CUDA图以优化重复推理 if not self.use_cuda_graph: return print(开始捕获CUDA图...) # 创建图 self.cuda_graph cuda.CUgraph() # 开始捕获 stream_handle self.stream.handle cuda.cuStreamBeginCapture(stream_handle, cuda.CUstreamCaptureMode.GLOBAL) # 执行一次推理这会被捕获到图中 dummy_input { name: np.random.randn(*info[shape]).astype(info[dtype]) for name, info in self.inputs.items() } _ self.infer(dummy_input) # 结束捕获 self.graph_executable cuda.cuStreamEndCapture(stream_handle) print(CUDA图捕获完成) def infer_optimized(self, input_data): 使用优化方法进行推理 if self.use_cuda_graph and self.cuda_graph is not None: # 使用CUDA图执行 # 注意输入数据形状必须和捕获时一致 self._update_inputs(input_data) cuda.cuGraphLaunch(self.graph_executable, self.stream) self.stream.synchronize() return self._get_outputs() else: # 回退到普通推理 return self.infer(input_data) def _update_inputs(self, input_data): 更新输入数据用于CUDA图 for input_info in self.inputs: name input_info[name] if name in input_data: data input_data[name] np.copyto(input_info[host], data.ravel()) cuda.memcpy_htod_async( input_info[device], input_info[host], self.stream ) def _get_outputs(self): 获取输出数据用于CUDA图 results {} for output_info in self.outputs: cuda.memcpy_dtoh_async( output_info[host], output_info[device], self.stream ) self.stream.synchronize() results[output_info[name]] output_info[host].reshape(output_info[shape]) return results def enable_tensor_core(self): 启用Tensor Core如果可用 if hasattr(self.context, set_optimization_profile_async): # 设置优化配置以使用Tensor Core profile_index 0 self.context.set_optimization_profile_async(profile_index, self.stream.handle) print(已启用Tensor Core优化) def set_compute_preference(self, preferencefastest): 设置计算偏好 # 这个需要在构建引擎时设置 # 这里只是展示概念 preferences { fastest: trt.BuilderPrecision.FASTEST, highest_precision: trt.BuilderPrecision.HIGHEST_PRECISION } if preference in preferences: print(f计算偏好设置为: {preference}) # 实际设置需要在builder config中完成CUDA图是个高级特性它能将一系列CUDA操作内存拷贝、kernel执行等打包成一个图然后重复执行这个图避免了每次推理时的调度开销。对于固定batch size和输入形状的推理场景CUDA图能带来显著的性能提升。但要注意一旦图被捕获输入输出的形状就不能变了。Tensor Core是英伟达Volta架构之后引入的专门用于矩阵乘加的硬件单元能极大加速深度学习计算。TensorRT会自动利用Tensor Core但你需要确保使用支持Tensor Core的数据类型如FP16和合适的形状。7. 实战踩坑与解决方案用了这么多年TensorRT我踩过的坑可能比很多人用过的次数还多。这里分享几个最常见的问题和解决方案希望能帮你少走弯路。问题一ONNX解析失败报错Unsupported ONNX opset version这是最常见的问题之一。TensorRT对ONNX算子集版本有要求太新或太老的版本都可能不支持。解决方案检查你的ONNX opset版本print(onnx.helper.printable_opset_version(model))TensorRT 10.x通常支持opset 11-15建议用opset 13导出ONNX时指定合适的版本torch.onnx.export(..., opset_version13)如果还是不行尝试用ONNX Simplifier简化模型python -m onnxsim input.onnx output.onnx问题二构建引擎时显存不足构建引擎需要大量显存特别是大模型或开启FP16/INT8时。解决方案减小max_workspace_size但不要小于256MB分批构建先构建部分层保存中间结果再构建剩余层使用builder.build_engine的serialization参数分步构建如果模型太大考虑模型剪枝或量化后再构建问题三推理结果与原始模型不一致精度问题通常由几个原因引起精度模式、算子实现差异、数据预处理不一致。解决方案先用FP32模式验证正确性再尝试FP16/INT8检查数据预处理归一化、通道顺序、数值范围使用trt.BuilderFlag.STRICT_TYPES强制类型一致性逐层对比输出找到第一个出现差异的层对于自定义算子检查TensorRT实现是否正确问题四动态形状推理性能下降动态形状虽然灵活但会有性能损失因为TensorRT需要为不同形状选择不同的kernel。解决方案限制动态范围不要设得太宽使用多个优化配置文件每个覆盖一个常用形状范围在服务端做输入标准化尽量使用固定尺寸对于batch维度动态使用trt.IExecutionContext的set_optimization_profile_async问题五多卡推理负载不均衡在多GPU服务器上如何充分利用所有GPU是个挑战。解决方案class MultiGPUInference: def __init__(self, engine_path, num_gpus4): self.engines [] self.contexts [] for gpu_id in range(num_gpus): cuda.init() device cuda.Device(gpu_id) context device.make_context() # 每个GPU加载一个引擎实例 with open(engine_path, rb) as f: runtime trt.Runtime(trt.Logger(trt.Logger.WARNING)) engine runtime.deserialize_cuda_engine(f.read()) self.engines.append(engine) self.contexts.append(engine.create_execution_context()) context.pop() self.current_gpu 0 self.num_gpus num_gpus def infer_round_robin(self, input_data): 轮询调度到不同GPU gpu_id self.current_gpu self.current_gpu (self.current_gpu 1) % self.num_gpus # 设置当前GPU上下文 device cuda.Device(gpu_id) context device.make_context() try: # 执行推理 # ... 推理代码 ... pass finally: context.pop() return result问题六内存泄漏长时间运行的服务可能会出现内存泄漏特别是频繁创建销毁TensorRT对象时。解决方案复用trt.Runtime和trt.ICudaEngine对象使用with语句确保资源正确释放定期检查GPU内存使用情况使用pycuda.driver的MemoryPool管理设备内存# 正确的资源管理示例 class SafeTensorRTInference: def __init__(self, engine_path): self.logger trt.Logger(trt.Logger.WARNING) # 使用with语句确保文件正确关闭 with open(engine_path, rb) as f: self.runtime trt.Runtime(self.logger) self.engine self.runtime.deserialize_cuda_engine(f.read()) self.context self.engine.create_execution_context() def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): # 清理资源 if hasattr(self, context): del self.context if hasattr(self, engine): del self.engine if hasattr(self, runtime): del self.runtime问题七批处理效率低小batch size无法充分利用GPU大batch size又增加延迟。解决方案使用动态批处理积累请求凑够一定batch size再推理根据请求延迟要求调整batch size使用TensorRT的IBuilderConfig设置最大batch size考虑使用流水线并行重叠数据传输和计算class DynamicBatcher: def __init__(self, max_batch_size32, timeout_ms10): self.max_batch_size max_batch_size self.timeout_ms timeout_ms self.batch_queue [] self.lock threading.Lock() def add_request(self, request): with self.lock: self.batch_queue.append(request) # 如果达到最大batch size或超时执行推理 if len(self.batch_queue) self.max_batch_size: return self._process_batch() return None def _process_batch(self): if not self.batch_queue: return [] # 合并多个请求为一个batch batch_data self._merge_requests(self.batch_queue) # 执行推理 results self.inference_engine.infer(batch_data) # 拆分结果返回给各个请求 split_results self._split_results(results) # 清空队列 self.batch_queue.clear() return split_results这些坑都是我亲身踩过的每个问题背后都可能浪费好几个小时的调试时间。希望这些经验能帮你避开这些陷阱。TensorRT虽然强大但毕竟是个复杂的系统遇到问题时不要慌仔细看错误信息从简单配置开始逐步调试总能找到解决方案的。