3、模型序列化与反序列化：.onnx转.trt，.uff转.trt，engine文件的保存与加载

好，咱们今天聊聊模型序列化。说白了，就是把训练好的模型，从一种格式变成另一种格式，最后存成TensorRT能直接用的engine文件。这个过程，我遇到过太多坑了。你想想看，模型在PyTorch里跑得好好的，一到TensorRT就报错，十有八九是序列化这步出了问题。

3.1 从ONNX到TRT：最常用的路径

我个人习惯，首选ONNX作为中间格式。为什么？因为它生态好，几乎主流框架都支持。把.onnx转成.trt，核心就两步：构建engine，然后序列化。

先看代码，这是我最常用的方式：

import tensorrt as trt

# 创建logger，TRT的日志全靠它
logger = trt.Logger(trt.Logger.WARNING)

# 创建builder
builder = trt.Builder(logger)

# 创建network，注意显式batch模式
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 创建parser，用来解析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))
        raise RuntimeError("ONNX解析失败")

# 配置builder
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 构建engine
serialized_engine = builder.build_serialized_network(network, config)

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

嗯，这里要注意。我在项目中遇到过，很多人直接调用builder.build_engine()，然后返回的是ICudaEngine对象。但新版本TRT推荐用build_serialized_network()，它直接返回序列化好的字节流，省去你手动序列化的步骤。

警告：ONNX解析失败，90%的原因是算子不支持。我建议你先用trtexec --onnx=model.onnx跑一下，看能不能成功。如果trtexec都报错，那代码肯定也不行。

3.2 从UFF到TRT：老路但还在用

UFF是TensorFlow的专用格式。说实话，现在用的人少了，但如果你维护的是老项目，可能还会碰到。我去年就接手过一个项目，模型是TensorFlow 1.x的，只能用UFF转。

流程类似，但parser不一样：

import tensorrt as trt
import uff

logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
network = builder.create_network()

# 先把TF模型转成UFF
uff_model = uff.from_tensorflow_frozen_model("frozen_model.pb", ["input_tensor"])

# 创建UFF parser
parser = trt.UffParser()
parser.register_input("input_tensor", (3, 224, 224), trt.UffInputOrder.NHWC)
parser.register_output("output_tensor")

# 解析UFF
parser.parse_buffer(uff_model, network)

# 构建engine
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)
engine = builder.build_engine(network, config)

# 序列化保存
with open("model.engine", "wb") as f:
    f.write(engine.serialize())

为什么UFF路径要单独提？因为它的输入顺序是NHWC，而ONNX默认是NCHW。我曾经在这上面栽过跟头，模型推理结果全是乱的，排查了半天才发现是输入顺序搞反了。

小技巧：如果你不确定输入顺序，可以用netron打开模型看一眼。UFF的输入节点会明确标注NHWC还是NCHW。

3.3 Engine文件的保存与加载

engine文件，说白了就是TensorRT优化后的“可执行文件”。它包含了网络结构、权重、以及针对你显卡的优化策略。所以，engine文件是跟显卡绑定的。你在一张RTX 3090上生成的engine，拿到RTX 4090上可能就跑不了。

保存很简单，上面已经展示了。加载呢？我一般这样写：

import tensorrt as trt

def load_engine(engine_path):
    logger = trt.Logger(trt.Logger.WARNING)
    with open(engine_path, "rb") as f:
        engine_data = f.read()
    
    runtime = trt.Runtime(logger)
    engine = runtime.deserialize_cuda_engine(engine_data)
    return engine

# 使用
engine = load_engine("model.trt")

这里有个细节。加载engine时，trt.Runtime会检查engine的兼容性。如果显卡驱动版本太低，或者CUDA版本不匹配，它会直接报错。我建议你在生产环境里，把engine的生成和加载分开。生成时记录下显卡型号和驱动版本，加载时做个校验。

核心要点：
  • ONNX转TRT：用OnnxParser，推荐build_serialized_network()
  • UFF转TRT：用UffParser，注意输入顺序NHWC
  • engine文件：跟显卡绑定，保存用serialize()，加载用deserialize_cuda_engine()

3.4 避坑指南：我踩过的那些坑

做模型序列化这几年，我总结了几条血泪教训：

  1. ONNX版本要匹配：TRT每个版本支持的ONNX opset版本不同。我建议你用ONNX 1.8以上，opset选13或更高。
  2. 动态shape要小心：如果你的模型输入是动态的（比如batch size可变），构建engine时要设置optimization_profile。否则推理时输入shape一变，engine就崩了。
  3. engine文件不要跨平台：Windows上生成的engine，Linux上不能用。反过来也一样。我有个同事把Windows上的engine拷到服务器上，折腾了一下午才发现问题。
  4. 序列化后的文件大小：如果模型权重很大，engine文件可能几个GB。保存时注意磁盘空间。我建议用gzip压缩一下，能省不少空间。
特别注意：千万不要在推理时频繁加载和卸载engine。每次deserialize_cuda_engine()都会重新分配显存，开销很大。正确的做法是：程序启动时加载一次，然后复用。

好了，关于模型序列化与反序列化，核心内容就这些。你想想看，其实流程不复杂，但细节特别多。我建议你动手试一遍，把ONNX转成TRT，再加载回来推理。遇到问题别慌，先看日志，再查算子支持列表。嗯，下一章咱们聊聊推理时的内存管理和显存优化，那个更有意思。

← 上一章 📖 返回目录 下一章 →