3、TensorRT基础：TensorRT简介、安装与配置、核心组件解析
3.1 TensorRT到底是什么？
说白了，TensorRT就是NVIDIA给自家GPU量身定做的一套推理加速引擎。它不是用来训练模型的，而是专门负责把训练好的模型跑得更快、更省显存。
我刚开始接触TensorRT时，有个误区——以为它就是个简单的模型转换工具。后来踩过坑才明白，它背后做的工作远比想象中复杂：层融合、精度校准、内核自动调优、内存复用……这些优化手段加在一起，推理速度能提升2-5倍，甚至更多。
你想想看，一个模型在PyTorch里跑一次推理要50ms，经过TensorRT优化后可能只要15ms。对于实时性要求高的场景，比如自动驾驶、直播特效，这差距就是能不能用的区别。
3.2 安装与配置——我踩过的坑
安装TensorRT其实不难，但版本匹配问题能让人抓狂。我个人习惯用tar包方式安装，因为灵活，想换版本就换。
3.2.1 环境要求
| 组件 | 推荐版本 | 备注 |
|---|---|---|
| CUDA | 11.x / 12.x | 必须与TensorRT版本对应 |
| cuDNN | 8.6+ | 建议用TensorRT配套的版本 |
| Python | 3.8 - 3.11 | 太新或太旧都可能出问题 |
| 操作系统 | Ubuntu 20.04 / 22.04 | Windows也能用，但坑更多 |
3.2.2 安装步骤（tar包方式）
# 1. 下载TensorRT tar包（以8.6.1为例）
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/secure/8.6.1/tars/TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz
# 2. 解压到指定目录
tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz -C /opt/
# 3. 配置环境变量（建议写到~/.bashrc里）
export LD_LIBRARY_PATH=/opt/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH
export PATH=/opt/TensorRT-8.6.1.6/bin:$PATH
# 4. 安装Python包
cd /opt/TensorRT-8.6.1.6/python
pip install tensorrt-8.6.1-cp39-none-linux_x86_64.whl
# 5. 验证安装
python -c "import tensorrt as trt; print(trt.__version__)"
3.3 核心组件解析——Builder、Engine、Context
这三个组件是TensorRT的骨架。理解它们的关系，你就掌握了TensorRT的80%。
我画了一张图来展示它们之间的协作关系：
3.3.1 Builder——模型的建筑师
Builder负责接收网络定义，然后根据你的配置（比如精度、工作空间大小、优化策略）生成一个最优的推理引擎。它只在构建阶段存在，推理时不需要它。
我记得第一次用Builder时，忘了设置最大工作空间，结果模型构建失败。后来养成习惯，每次都会显式设置：
import tensorrt as trt
# 创建Builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
# 创建网络定义（显式batch模式）
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
# 配置优化参数
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) # 1GB工作空间
config.set_flag(trt.BuilderFlag.FP16) # 启用FP16推理
# 设置动态shape范围（如果模型支持）
profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 3, 224, 224), (4, 3, 224, 224), (8, 3, 224, 224))
config.add_optimization_profile(profile)
3.3.2 Engine——序列化的推理引擎
Engine是Builder产出的最终产物。它是一个高度优化的、序列化后的模型表示，可以保存为.plan文件。下次加载时直接反序列化就能用，不需要重新构建。
这里有个关键点：Engine是平台相关的。你在A100上构建的Engine，不能直接拿到T4上用。为什么？因为不同GPU的架构不同，TensorRT针对特定架构做了内核调优。
# 构建Engine并序列化
serialized_engine = builder.build_serialized_network(network, config)
with open("resnet50_fp16.plan", "wb") as f:
f.write(serialized_engine)
# 从文件加载Engine
with open("resnet50_fp16.plan", "rb") as f:
serialized_engine = f.read()
runtime = trt.Runtime(logger)
engine = runtime.deserialize_cuda_engine(serialized_engine)
3.3.3 Context——执行推理的工人
Context是实际干活的那个。它管理着推理过程中的所有状态：输入输出缓冲区的绑定、CUDA流的分配、内核的执行。
一个Engine可以创建多个Context，实现多路并发推理。但要注意，每个Context都有自己的状态，不是线程安全的。
# 创建Context
context = engine.create_execution_context()
# 准备输入输出缓冲区
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
# 分配GPU内存
input_buf = cuda.mem_alloc(1 * 3 * 224 * 224 * 4) # float32
output_buf = cuda.mem_alloc(1 * 1000 * 4)
# 设置输入shape（动态shape时需要）
context.set_binding_shape(0, (1, 3, 224, 224))
# 执行推理
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
cuda.memcpy_htod(input_buf, input_data)
context.execute_v2([int(input_buf), int(output_buf)])
# 取回结果
output_data = np.empty(1000, dtype=np.float32)
cuda.memcpy_dtoh(output_data, output_buf)
print(output_data[:5])
3.4 三者的协作流程
总结一下，整个流程是这样的：
- 构建阶段： Builder读取网络定义（ONNX等），根据配置进行优化，输出Engine
- 序列化阶段： Engine保存为.plan文件，方便分发和部署
- 反序列化阶段： 加载.plan文件，通过Runtime创建Engine对象
- 推理阶段： 从Engine创建Context，绑定输入输出，执行推理
嗯，这里要注意：Builder和Runtime是两回事。Builder负责构建，Runtime负责加载。生产环境中，你只需要Runtime就够了，因为Engine已经构建好了。
我个人习惯把构建和推理分开写两个脚本。构建脚本在开发机上跑一次，生成.plan文件。推理脚本部署到生产环境，只加载Engine执行推理。这样既保证了性能，又简化了部署流程。
公众号:蓝海资料掘金营，微信deep3321