4. TensorRT基础入门：核心组件与优化策略

说实话，TensorRT这玩意儿刚接触时挺唬人的。一堆新名词：Builder、Engine、Context、ONNX解析器……我当时也懵过。但用久了你会发现，核心逻辑其实很清晰。今天我就带你把这些组件拆开揉碎，讲明白。

4.1 三大核心组件：Builder、Engine、Context

这三个东西，说白了就是“搭积木”的三个阶段。

4.1.1 Builder（构建器）

Builder是TensorRT的“总设计师”。它负责读取模型、设置参数、选择优化策略。我个人习惯把Builder想象成一个项目经理——它不干活，但决定怎么干。

创建Builder的代码很简单：

import tensorrt as trt

# 创建Builder
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
# 设置工作空间大小（单位：字节）
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

嗯，这里要注意：工作空间设太大浪费显存，设太小可能优化失败。我一般从1GB开始试，不够再加。

4.1.2 Engine（引擎）

Engine是优化后的“可执行文件”。Builder把模型优化完，就生成一个Engine。你可以把它序列化保存到硬盘，下次直接加载，省去重新优化的时间。

我在项目中遇到过一个问题：Engine是硬件绑定的。你在A100上生成的Engine，拿到T4上可能跑不了。为什么？因为TensorRT会根据具体GPU的架构做针对性优化。

# 序列化Engine
serialized_engine = builder.build_serialized_network(network, config)
with open("model.engine", "wb") as f:
    f.write(serialized_engine)

# 反序列化加载
with open("model.engine", "rb") as f:
    engine_data = f.read()
runtime = trt.Runtime(trt.Logger(trt.Logger.WARNING))
engine = runtime.deserialize_cuda_engine(engine_data)

4.1.3 Context（执行上下文）

Context是真正干活的。它管理着推理时的内存、CUDA流、中间结果。一个Engine可以创建多个Context，实现多路并发推理。

# 创建Context
context = engine.create_execution_context()

# 准备输入输出缓冲区
input_idx = engine.get_binding_index("input")
output_idx = engine.get_binding_index("output")

# 分配GPU内存
input_buffer = cuda.mem_alloc(input_size * input_dtype.itemsize)
output_buffer = cuda.mem_alloc(output_size * output_dtype.itemsize)

# 执行推理
cuda.memcpy_htod(input_buffer, input_data)
context.execute_v2([int(input_buffer), int(output_buffer)])
cuda.memcpy_dtoh(output_data, output_buffer)

4.2 ONNX解析器：连接PyTorch与TensorRT的桥梁

你想想看，我们训练模型用的是PyTorch、TensorFlow，但TensorRT不认识这些框架。怎么办？ONNX就是中间人。

ONNX解析器的作用很简单：把ONNX模型文件读进来，转换成TensorRT内部的Network定义。

# 创建ONNX解析器
parser = trt.OnnxParser(network, logger)

# 读取ONNX文件
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))

我曾经踩过一个坑：ONNX模型里有些算子TensorRT不支持。比如某些自定义的激活函数。解决办法有两个：要么改模型，用TensorRT支持的算子替代；要么用Plugin自己写。

4.3 网络定义：手动搭建还是自动解析？

TensorRT支持两种方式定义网络：

• 通过ONNX解析器自动导入：省事，适合大多数场景
• 手动用API搭建：灵活，适合特殊需求

手动搭建的例子：

network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 添加输入层
input_tensor = network.add_input("input", trt.float32, (1, 3, 224, 224))

# 添加卷积层
conv = network.add_convolution(input_tensor, 64, (3, 3), conv_weight, conv_bias)
conv.stride = (1, 1)
conv.padding = (1, 1)

# 添加ReLU
relu = network.add_activation(conv.get_output(0), trt.ActivationType.RELU)

# 标记输出
network.mark_output(relu.get_output(0))

说实话，手动搭建太累了。除非你要做非常特殊的网络结构，否则我建议直接用ONNX解析器。

4.4 优化策略：TensorRT凭什么这么快？

TensorRT的优化策略，说白了就是三板斧：

优化策略	说明	效果
层融合（Layer Fusion）	把多个连续的小算子合并成一个	减少Kernel启动开销
精度校准（INT8/FP16）	用更低精度替代FP32	显存减半，速度翻倍
内核自动调优	为每个算子选择最优的CUDA Kernel	充分利用硬件特性
内存复用	复用中间张量的内存空间	降低显存占用

我印象最深的是层融合。你看一个简单的Conv+BN+ReLU，在PyTorch里是三个独立的操作。TensorRT会把它融合成一个CBR算子，省掉了中间结果的读写。速度能快30%以上。

4.5 知识体系总览

下面这张图，是我自己总结的TensorRT核心流程。你看一遍，基本就明白整个链路了。

这张图把整个流程串起来了。你从左边ONNX模型开始，经过解析器变成网络定义，Builder带着优化策略生成Engine，最后Context执行推理。每一步都有它的职责，缺一不可。

4.6 实战建议

最后，给你几个我踩坑换来的建议：

• 先用FP16试试：大多数模型FP16精度损失很小，但速度提升明显。INT8需要校准数据集，麻烦一些。
• 保存Engine文件：构建过程很慢，尤其是大模型。一次构建，多次加载，能省很多时间。
• 注意动态形状：如果你的输入尺寸会变，记得设置动态形状优化。否则固定尺寸的Engine会报错。
• 日志级别别设太高：我习惯用WARNING级别。INFO太多，ERROR又可能漏掉重要信息。