模型导出与格式标准：ONNX、SavedModel/PB 与 TensorRT 专属格式

模型训练好了，只是第一步。真正要让它跑起来，还得过「导出」这一关。我见过不少团队，模型在 PyTorch 里精度漂亮得很，一导出就崩，一部署就炸。说白了，格式转换这件事，坑比想象中多。

今天咱们就聊聊三种主流格式：ONNX、SavedModel/PB、以及 TensorRT 的 .engine。我会结合自己踩过的坑，把每个格式的要点讲清楚。

ONNX：模型界的「通用语言」

ONNX 全称 Open Neural Network Exchange，微软和 Meta 搞出来的。它的目标很简单——让模型能在不同框架间自由迁移。我个人觉得，ONNX 最大的价值不是转换本身，而是它提供了一个「中间表示」，让你能脱离原始框架做推理优化。

ONNX 的核心结构

一个 ONNX 模型，本质上是个 protobuf 文件。里面包含三样东西：

  • 计算图（Graph）：节点是算子，边是张量流动
  • 权重参数：以 TensorProto 形式存储
  • 元数据：模型版本、算子集版本、生产者信息等

嗯，这里要注意：ONNX 的算子集是分版本的。你导出的模型用了 opset 13，但推理引擎只支持 opset 11，那就得做算子降级。我在项目中遇到过，一个简单的 Resize 操作，不同 opset 下的行为完全不同，搞不好就出边界问题。

PyTorch 导出 ONNX 的实战

PyTorch 导出 ONNX 主要靠 torch.onnx.export。看起来就一行代码，但参数调起来挺讲究。

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    export_params=True,
    opset_version=13,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里我重点说几个参数：

  • dynamic_axes：允许动态 batch 大小。如果你不设，导出的模型就固定 batch=1，部署时改不了。我曾经因为这个被运维同事追着骂——模型只能一次推理一张图，线上 QPS 根本扛不住。
  • do_constant_folding：把常量计算提前折叠掉。比如权重和偏置的加法，导出时就算好，推理时少一步操作。能省一点是一点。
  • opset_version：建议用 13 或更高。太老的版本有些算子不支持，太新的版本推理引擎可能不兼容。
注意：PyTorch 的 torch.onnx.export 走的是 trace 模式，不是 script 模式。这意味着它只记录你给的那个 dummy_input 的执行路径。如果你的模型里有条件分支（if-else），没走到的分支就不会被导出。我见过有人模型里有个 if training: 的分支，导出时 training=False，结果推理时那个分支的权重根本没导进去，精度直接崩了。

ONNX 的验证与调试

导出完别急着用，先验证一下：

import onnx
import onnxruntime as ort

# 检查模型结构
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)

# 用 onnxruntime 跑一遍
ort_session = ort.InferenceSession("resnet18.onnx")
outputs = ort_session.run(None, {"input": dummy_input.numpy()})
print(outputs[0].shape)

你想想看，如果这一步输出和 PyTorch 原始输出对不上，那后面部署全是白费功夫。我习惯把 PyTorch 和 ONNX 的输出做逐元素对比，误差超过 1e-5 就要警惕了。

TensorFlow 的 SavedModel 与 PB

TensorFlow 的模型格式有点乱，早期是 checkpoint，后来是 frozen graph（.pb），再后来是 SavedModel。现在官方推荐的是 SavedModel，但很多老项目还在用 .pb。

SavedModel 是什么？

SavedModel 是一个目录，里面包含：

  • saved_model.pb：模型的计算图定义
  • variables/：权重参数文件
  • assets/：额外的资源文件（比如词汇表）

说白了，SavedModel 比 .pb 更「完整」。它保留了模型的签名信息，你部署时不用自己拼输入输出张量的名字。我刚开始用 TensorFlow Serving 时，就因为不知道输入张量叫什么名字，折腾了一下午。

导出 SavedModel 的方法

import tensorflow as tf

# 假设你有一个 Keras 模型
model = tf.keras.applications.MobileNetV2(weights='imagenet')

# 导出 SavedModel
tf.saved_model.save(model, "mobilenetv2_savedmodel")

# 也可以用具体函数签名导出
@tf.function(input_signature=[tf.TensorSpec(shape=[None, 224, 224, 3], dtype=tf.float32)])
def infer(input_tensor):
    return model(input_tensor)

tf.saved_model.save(model, "mobilenetv2_signed", signatures={"serving_default": infer})

我个人建议用第二种方式，显式指定签名。这样部署时接口清晰，不会出现「模型能加载但不知道怎么调用」的尴尬。

PB 文件（冻结图）

PB 文件是把计算图和权重打包成一个文件。它没有 SavedModel 那么灵活，但胜在轻量，适合嵌入式设备。

# 从 checkpoint 导出 frozen graph
import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

# 加载模型
model = tf.keras.applications.MobileNetV2(weights='imagenet')
full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(
    tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))

# 冻结
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()

# 保存为 .pb
tf.io.write_graph(frozen_func.graph, "./", "mobilenetv2.pb", as_text=False)
小技巧：PB 文件可以用 as_text=True 保存为文本格式，方便查看计算图结构。但部署时一定要用二进制格式，文件小、加载快。

TensorRT 的专属格式：.engine

TensorRT 是 NVIDIA 的推理优化引擎。它不直接吃 ONNX 或 SavedModel，而是把它们编译成自己的 .engine 文件。这个编译过程，说白了就是把计算图做「手术」——算子融合、精度校准、内存复用，一套组合拳下来，推理速度能快好几倍。

从 ONNX 到 TensorRT Engine

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)

# 解析 ONNX 模型
with open("resnet18.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for error in range(parser.num_errors):
            print(parser.get_error(error))

# 构建配置
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("resnet18.engine", "wb") as f:
    f.write(serialized_engine)

这里有个坑：TensorRT 的 builder 非常吃显存。我曾在 8GB 显存的卡上编译一个 ResNet-50，直接 OOM 了。后来发现是 workspace 设得太大。你想想看，编译时 workspace 是给 TensorRT 做算子搜索用的，不是推理时用的，设个 1GB 通常就够了。

TensorRT 的精度校准

TensorRT 支持 FP32、FP16、INT8 三种精度。FP16 基本无损，INT8 就需要做校准了——用一批真实数据跑一遍，统计每层激活值的分布，然后决定量化参数。

# INT8 校准器示例
class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, data_loader, cache_file):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.data_loader = data_loader
        self.cache_file = cache_file
        self.batch_size = 32
        self.current_index = 0
        
    def get_batch_size(self):
        return self.batch_size
    
    def get_batch(self, names):
        if self.current_index + self.batch_size > len(self.data_loader):
            return None
        batch = self.data_loader[self.current_index:self.current_index + self.batch_size]
        self.current_index += self.batch_size
        return [batch.numpy()]
    
    def read_calibration_cache(self):
        if os.path.exists(self.cache_file):
            with open(self.cache_file, "rb") as f:
                return f.read()
    
    def write_calibration_cache(self, cache):
        with open(self.cache_file, "wb") as f:
            f.write(cache)
重要提醒：INT8 校准用的数据一定要和实际推理数据分布一致。我曾经用 ImageNet 的验证集做校准，结果部署到工业场景（全是黑白文档扫描图），精度掉了 5 个点。后来换了真实场景数据重新校准，才恢复正常。

三种格式的对比

特性 ONNX SavedModel/PB TensorRT .engine
跨框架支持 强（PyTorch/TF/…） 弱（主要 TF 生态） 弱（仅 NVIDIA GPU）
推理速度 中等 中等 极快（算子融合+量化）
可读性 好（protobuf 可解析） 一般 差（二进制序列化）
部署灵活性 高（多种 runtime） 中（TF Serving/Java） 低（仅 TensorRT）
典型场景 模型中转、跨平台 TF 生态部署 高性能 GPU 推理

我个人习惯的流程是：PyTorch 训练 → 导出 ONNX → 用 ONNX Runtime 验证精度 → 再转 TensorRT engine 做高性能部署。每一步都做精度对齐，确保不出幺蛾子。

嗯，模型导出这件事，看起来是「导出」两个字，实际上涉及算子兼容性、精度对齐、性能优化一堆事。你想想看，一个模型从训练到上线，中间要过多少道手？每一道手都可能引入误差。所以我的建议是：每转一次格式，就做一次精度验证。别等到线上出问题了再回头查，那时候就晚了。

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