1. TensorRT 简介
TensorRT 是 NVIDIA 推出的高性能深度学习推理 SDK，核心目标是将训练好的模型转化为能够在 NVIDIA GPU 上高效执行的推理引擎。与基于 PyTorch、TensorFlow 的原生推理相比，TensorRT 通过层融合（Layer Fusion）、精度校准（Precision Calibration）和张量内存优化等手段，通常能带来 2~10 倍 的推理加速，同时显著降低延迟和显存占用。
TensorRT 支持的主流平台覆盖：
- 数据中心级 GPU：A100、H100、A10、T4 等
- 工作站/服务器：x86_64 + NVIDIA GPU
- 边缘设备：NVIDIA Jetson 系列（ARM64 + integrated GPU），支持 Orin、Xavier 等
对于高并发在线服务或资源受限的边缘场景，TensorRT 几乎是 NVIDIA 生态下的必选项。
2. 核心工作流四阶段
TensorRT 将模型部署划分为四个阶段，理解这一流水线是调优的基础。
Phase 1: Parse（解析）
TensorRT 支持从 ONNX、UFF、Caffe 等格式解析网络结构。当前最推荐的输入格式是 ONNX，绝大多数框架都原生支持导出，且 TensorRT 对 ONNX 的算子覆盖最为完善。
import tensorrt as trt
logger = trt.Logger(trt.Logger.INFO)
builder = trt.Builder(logger)
network = builder.create_network(
1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, logger)
with open("model.onnx", "rb") as f:
if not parser.parse(f.read()):
for error in range(parser.num_errors):
print(parser.get_error(error))
Phase 2: Build（构建）
Builder 负责执行图优化与内核搜索。通过 IBuilderConfig 可以精确控制优化行为：
config = builder.create_builder_config()
config.max_workspace_size = 4 * 1024 * 1024 * 1024 # 4 GB
config.set_flag(trt.BuilderFlag.FP16)
# 可选：启用 INT8（需配合 Calibrator）
# config.set_flag(trt.BuilderFlag.INT8)
# config.int8_calibrator = MyCalibrator(...)
# Jetson 平台可启用 DLA
# config.set_flag(trt.BuilderFlag.GPU_FALLBACK)
# config.default_device_type = trt.DeviceType.DLA
关键配置项：
| 配置项 | 说明 |
|---|---|
FP16 | 半精度推理，几乎无损，速度提升明显 |
INT8 | 整型推理，需要校准，加速最高，适合吞吐敏感型应用 |
max_workspace_size | Builder 可用的临时显存上限，越大越有利于找到更快 kernel |
DLA | Deep Learning Accelerator，Jetson 专用硬件加速器 |
Phase 3: Engine（引擎）
构建完成后，TensorRT 生成一个高度优化的 ICudaEngine 对象。该对象包含经过层融合、内核自动调优后的执行图，通常直接将其序列化到磁盘，避免每次启动都重新构建。
engine = builder.build_engine(network, config)
with open("model.engine", "wb") as f:
f.write(engine.serialize())
Phase 4: Inference（推理）
运行时反序列化引擎，并通过 IExecutionContext 执行推理：
runtime = trt.Runtime(logger)
with open("model.engine", "rb") as f:
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()
# 绑定输入输出显存后执行 context.execute_async_v2(bindings, stream)
3. 网络层优化机制
TensorRT 的加速并非来源于某个单一魔法，而是多个底层优化的叠加。
层融合（Layer Fusion）
Builder 会自动识别图中相邻且可合并的算子，将其融合为单一 CUDA kernel。典型例子如 Conv + Bias + ReLU，融合后消除了中间结果写回显存的开销，减少 kernel launch 次数。
张量内存优化
TensorRT 采用 显存复用（Memory Reuse） 策略：在引擎构建期计算每个张量的生命周期，不同层的中间张量可以共用同一段显存，显著降低峰值显存占用。
精度校准
- FP16：利用 Tensor Cores 加速，无需额外数据，是最基础的精度降级策略。
- INT8：需要逐层统计激活分布，通过校准器确定最佳量化参数。TensorRT 提供
IInt8EntropyCalibrator2（推荐用于 CNN）和IInt8MinMaxCalibrator（推荐用于 Transformer/BERT）。
class MyCalibrator(trt.IInt8EntropyCalibrator2):
def __init__(self, data_loader):
super().__init__()
self.data_loader = data_loader
self.batch_size = 32
self.cache_file = "calib.cache"
def get_batch_size(self):
return self.batch_size
def get_batch(self, names):
try:
batch = next(self.iter)
return [batch.ctypes.data]
except StopIteration:
return None
def read_calibration_cache(self):
try:
with open(self.cache_file, "rb") as f:
return f.read()
except FileNotFoundError:
return None
def write_calibration_cache(self, cache):
with open(self.cache_file, "wb") as f:
f.write(cache)
动态 Shape（Dynamic Shapes）
当输入尺寸不固定时（如可变 batch、可变图像分辨率），可以通过配置 Optimization Profile 告诉 TensorRT 预期的 shape 范围：
profile = builder.create_optimization_profile()
profile.set_shape("input", min=(1, 3, 224, 224),
opt=(8, 3, 224, 224),
max=(32, 3, 224, 224))
config.add_optimization_profile(profile)
实际推理前，通过 context.set_binding_shape 或 context.set_input_shape（TRT 8.6+）设定当前批次的确切尺寸。
4. 高级特性
自定义 Plugin（IPluginV2）
当网络中出现 TensorRT 不支持的算子时，你可以编写自定义 Plugin。自 TensorRT 8 起推荐使用 IPluginV2DynamicExt 接口：
class MyCustomPlugin : public nvinfer1::IPluginV2DynamicExt {
public:
// 返回输出维度，支持动态 shape
nvinfer1::DimsExprs getOutputDimensions(
int outputIndex, const nvinfer1::DimsExprs* inputs,
int nbInputs, nvinfer1::IExprBuilder& exprBuilder) override;
// 配置 CUDA kernel launch 参数
int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
const nvinfer1::PluginTensorDesc* outputDesc,
const void* const* inputs, void* const* outputs,
void* workspace, cudaStream_t stream) override;
// 序列化/反序列化插件参数
size_t getSerializationSize() const override;
void serialize(void* buffer) const override;
};
编写 Plugin 时应注意：
- 精确实现
supportsFormatCombination，声明支持的 data type 与 format enqueue中尽量复用已编译的 CUDA kernel，避免动态编译- 在 Python 侧通过
PLUGIN_LIBRARY导入动态链接库并注册 creator
Explicit Batch vs Implicit Batch
- Implicit Batch（旧模式）：batch 维度由 TensorRT 隐式处理，不支持动态 batch。TensorRT 7.x 及之前版本的默认模式。
- Explicit Batch（推荐）：batch 作为显式维度出现在张量 shape 中，ONNX 导出天然匹配此模式。自 TensorRT 7 起即可启用，TensorRT 8+ 为推荐方式。
TensorRT 8.x 与 10.x API 演进
TensorRT 10 是一次较大重构，移除了大量已弃用接口：
| 特性 | TensorRT 8.x | TensorRT 10.x |
|---|---|---|
| 输入 shape 设置 | set_binding_shape | set_input_shape |
| 获取输入名 | get_binding_name | get_io_tensor_name |
| Engine 绑定 | execute_async_v2(bindings) | execute_async_v3（无需 bindings 数组） |
| DLA 支持 | 完整支持 | 在部分 API 上有调整 |
迁移时建议使用 NVIDIA 提供的 polygraphy 或 trtexec 先做兼容性验证。
5. 序列化与部署策略
将构建好的 engine 序列化到磁盘是典型的生产部署做法：
# 构建期一次性执行
def build_engine(onnx_path, engine_path):
...
engine = builder.build_engine(network, config)
with open(engine_path, "wb") as f:
f.write(engine.serialize())
需要注意：
- 版本锁定：TensorRT engine 对 TRT 版本、CUDA 版本、GPU 架构严格绑定。A100 上构建的 engine 无法在 T4 上直接运行，跨版本也无法加载。
- 多 Engine 策略：如果业务场景包含差异较大的 batch size（如 1 vs 64），可以构建多个 engine，每个 engine 在各自 optimized shape 附近性能最佳。
6. Python 完整流程示例
以下脚本展示从 ONNX 导出到推理 benchmark 的全链路：
import tensorrt as trt
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit
def build_engine_from_onnx(onnx_file, engine_file, fp16=True):
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network(
1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, logger)
with open(onnx_file, "rb") as f:
parser.parse(f.read())
config = builder.create_builder_config()
config.max_workspace_size = 2 << 30
if fp16:
config.set_flag(trt.BuilderFlag.FP16)
# 动态 shape profile（可选）
profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 3, 224, 224), (4, 3, 224, 224), (16, 3, 224, 224))
config.add_optimization_profile(profile)
engine = builder.build_engine(network, config)
with open(engine_file, "wb") as f:
f.write(engine.serialize())
return engine
def allocate_buffers(engine, batch_size):
inputs, outputs, bindings = [], [], []
stream = cuda.Stream()
for i in range(engine.num_io_tensors):
name = engine.get_tensor_name(i)
mode = engine.get_tensor_mode(name)
shape = (batch_size,) + tuple(engine.get_tensor_shape(name)[1:])
dtype = trt.nptype(engine.get_tensor_dtype(name))
size = trt.volume(shape)
host_mem = cuda.pagelocked_empty(size, dtype)
device_mem = cuda.mem_alloc(host_mem.nbytes)
bindings.append(int(device_mem))
if mode == trt.TensorIOMode.INPUT:
inputs.append({"name": name, "host": host_mem, "device": device_mem})
else:
outputs.append({"name": name, "host": host_mem, "device": device_mem})
return inputs, outputs, bindings, stream
def infer(engine, input_data):
batch_size = input_data.shape[0]
context = engine.create_execution_context()
context.set_input_shape("input", input_data.shape)
inputs, outputs, bindings, stream = allocate_buffers(engine, batch_size)
np.copyto(inputs[0]["host"], input_data.ravel())
cuda.memcpy_htod_async(inputs[0]["device"], inputs[0]["host"], stream)
context.execute_async_v3(stream_handle=stream.handle)
cuda.memcpy_dtoh_async(outputs[0]["host"], outputs[0]["device"], stream)
stream.synchronize()
return outputs[0]["host"].reshape(batch_size, -1)
# trtexec 命令行快速验证
# trtexec --onnx=model.onnx --saveEngine=model.engine --fp16 --minShapes=input:1x3x224x224 \
# --optShapes=input:4x3x224x224 --maxShapes=input:16x3x224x224
Benchmark 时建议采用 trtexec 工具，它能自动输出各层的延迟分解和吞吐：
trtexec --loadEngine=model.engine --shapes=input:16x3x224x224 --warmUp=500 --duration=30
7. 常见陷阱与对策
| 问题 | 原因 | 对策 |
|---|---|---|
Unsupported operator | TensorRT 未实现该算子 | 1) 替换为等价支持算子组合；2) 编写 Custom Plugin |
| 构建时 OOM | workspace 过大或模型过大 | 分阶段构建、减小 workspace、启用显存池隔离 |
| INT8 精度下降严重 | 校准数据集分布与真实数据差异大 | 使用>=500张真实业务样本做校准；选择正确的 Calibrator |
| 动态 shape 延迟高 | 未命中 optimized profile | 确保运行时 shape 落在 opt 附近；为常用 shape 单独建 engine |
| Engine 加载失败 | 跨版本或跨 GPU 架构 | 在目标硬件上重新构建 engine；做好 CI 中的版本一致性校验 |
| Plugin 找不到 | 未注册或未加载动态库 | Python 中 ctypes.CDLL 加载 .so；确认 register_creator 已调用 |
总结
TensorRT 是 NVIDIA 推理栈的核心调速器。掌握其 “Parse → Build → Engine → Infer” 四阶段流水线后，重点投入方向应聚焦于：
- 精度与速度的权衡：优先尝试 FP16，再在吞吐敏感路径上引入 INT8；
- 动态 Shape 的正确配置：通过 optimization profile 保证 engine 在常用 shape 上达到峰值性能；
- 部署稳定性：严格保证构建环境与运行环境的一致性，将 engine 构建纳入 CI/CD 流程。
在高性能推理场景中，TensorRT 与 CUDA Graph、Multi-Stream 等技术打组合拳，往往能将 GPU 利用率推向极限。后续可以进一步探索 Triton Inference Server + TensorRT backend 的高并发服务化部署方案。
继续阅读
探索更多技术文章
浏览归档，发现更多关于系统设计、工具链和工程实践的内容。