4. TensorRT Python API入门：创建Builder与Config、构建Network定义、设置优化策略、生成Engine文件

好，咱们今天来聊聊TensorRT的Python API。说实话，很多同学一上来就急着调inference，结果连Engine怎么生成的都搞不清楚。我个人习惯是，先把这几个核心对象摸透：Builder、Config、Network，最后才是Engine。你想想看，盖房子总得先有图纸、有施工队、有监理吧？TensorRT里也是这个道理。

4.1 创建Builder与Config

Builder是TensorRT的“总指挥”。它负责协调整个模型优化和编译的过程。我刚开始用的时候，以为Builder就是个工厂类，直接new一下就行。其实没那么简单。

import tensorrt as trt

# 创建Builder - 这是入口
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 创建Config - 这是配置中心
config = builder.create_builder_config()

# 设置工作空间大小（显存上限）
# 我曾经遇到过显存爆炸的问题，就是因为没设这个
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

这里有个坑：trt.Logger的级别。我建议开发阶段用INFO，生产环境用WARNING。为什么？因为INFO会打印每一层的详细信息，调试时很有用，但线上跑起来日志量太大。

4.2 构建Network定义

Network定义就是你的模型“图纸”。TensorRT支持两种方式构建网络：一种是直接解析ONNX或Caffe模型，另一种是手动搭建网络层。咱们先讲最常用的——从ONNX解析。

# 创建Network对象
# 显式指定batch维度，这样后续可以动态调整batch size
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 加载ONNX模型
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))
        raise RuntimeError("ONNX解析失败")

print(f"网络输入数量: {network.num_inputs}")
print(f"网络输出数量: {network.num_outputs}")

嗯，这里要注意。ONNX解析失败是家常便饭。我遇到过最离谱的一次，是因为模型里有个自定义算子，ONNX标准里没有。后来我手动用Plugin实现了那个算子才搞定。

4.3 设置优化策略

优化策略说白了就是告诉TensorRT：“嘿，你帮我试试这几种配置，看哪个跑得最快。”Config里有个optimization_profile的概念，专门干这个的。

# 创建优化profile
profile = builder.create_optimization_profile()

# 设置输入张量的形状范围
# 格式: (最小形状, 最优形状, 最大形状)
input_name = network.get_input(0).name
profile.set_shape(input_name, 
                  min_shape=(1, 3, 224, 224),
                  opt_shape=(8, 3, 224, 224),
                  max_shape=(32, 3, 224, 224))

config.add_optimization_profile(profile)

# 设置精度模式
# 我建议先跑FP32验证正确性，再切FP16或INT8
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)
    print("支持FP16加速，已启用")

# 设置严格类型约束（默认关闭）
# config.set_flag(trt.BuilderFlag.STRICT_TYPES)

你可能会问：为什么要设三个形状？说白了，TensorRT需要知道你的输入范围，它才能针对性地做内存规划和算子选择。如果你只给一个固定形状，那动态batch就用不了了。

参数	含义	我的建议
min_shape	最小输入尺寸	设为实际业务的最小值
opt_shape	最常见尺寸	设为batch size的典型值
max_shape	最大输入尺寸	不要超过显存上限

4.4 生成Engine文件

终于到这一步了。Engine就是TensorRT编译后的“可执行文件”。它包含了优化后的网络结构、权重、以及CUDA kernel代码。说白了，Engine就是你的模型“成品”。

# 序列化Engine
# 这一步可能会比较慢，取决于模型大小
engine_bytes = builder.build_serialized_network(network, config)

if engine_bytes is None:
    raise RuntimeError("Engine构建失败")

# 保存到文件
with open("model.engine", "wb") as f:
    f.write(engine_bytes)

print(f"Engine文件大小: {len(engine_bytes) / 1024 / 1024:.2f} MB")

# 或者直接创建runtime可用的Engine
# runtime = trt.Runtime(logger)
# engine = runtime.deserialize_cuda_engine(engine_bytes)

这里有个关键点：build_serialized_network返回的是序列化的字节流，不是直接可用的Engine对象。你需要用Runtime来反序列化。我习惯把Engine保存成文件，这样下次启动时直接加载，省去重新编译的时间。

4.5 完整流程示例

把上面的步骤串起来，就是一个完整的Engine生成流程。我一般会封装成一个函数，方便复用。

def build_engine(onnx_path, engine_path, batch_size=8):
    """
    从ONNX生成TensorRT Engine
    
    Args:
        onnx_path: ONNX模型路径
        engine_path: 输出Engine路径
        batch_size: 默认batch size
    """
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    config = builder.create_builder_config()
    
    # 设置工作空间
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 2 << 30)  # 2GB
    
    # 创建网络
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )
    
    # 解析ONNX
    parser = trt.OnnxParser(network, logger)
    with open(onnx_path, "rb") as f:
        if not parser.parse(f.read()):
            for i in range(parser.num_errors):
                print(parser.get_error(i))
            return None
    
    # 设置优化profile
    profile = builder.create_optimization_profile()
    input_name = network.get_input(0).name
    profile.set_shape(input_name,
                      (1, 3, 224, 224),
                      (batch_size, 3, 224, 224),
                      (batch_size * 4, 3, 224, 224))
    config.add_optimization_profile(profile)
    
    # 启用FP16（如果支持）
    if builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    
    # 构建Engine
    print("开始构建Engine，这可能需要几分钟...")
    engine_bytes = builder.build_serialized_network(network, config)
    
    if engine_bytes:
        with open(engine_path, "wb") as f:
            f.write(engine_bytes)
        print(f"Engine已保存到: {engine_path}")
        return engine_bytes
    else:
        print("Engine构建失败")
        return None

# 使用示例
build_engine("resnet50.onnx", "resnet50.engine", batch_size=8)

你看，整个流程其实不复杂。但每个环节都有它的门道。我个人觉得，最容易被忽视的就是优化profile的设置。很多人随便填个形状就完事了，结果线上性能差一大截。

好了，这一章的内容就到这里。下一章咱们会讲如何加载Engine并执行推理，到时候你会看到，前面这些准备工作都是值得的。