模型导出：从训练框架到部署的桥梁

模型训练完了，怎么把它送到推理引擎手里？这一步，说白了就是模型导出。我做了这么多年模型部署，见过太多人在这一步翻车。今天咱们就聊聊 PyTorch 和 TensorFlow 这两大主流框架的导出方式。

你想想看，训练框架里跑得好好的模型，到了部署环境可能连 Python 都没有。所以我们需要一个中间表示——既保留模型的计算图，又去掉训练相关的包袱。

PyTorch 模型导出：TorchScript 与 ONNX

PyTorch 的导出方式，我个人习惯分成两条路：TorchScript 和 ONNX。前者是 PyTorch 自家的，后者是跨框架的。

torch.jit.trace 与 torch.jit.script

先说 torch.jit.trace。它的原理很简单：给模型一个输入，跑一遍前向，把执行路径记录下来。嗯，这里要注意——它只记录实际执行到的代码路径。

核心区别：trace 是追踪执行路径，script 是编译整个模型代码。

# trace 方式
import torch

class MyModel(torch.nn.Module):
    def forward(self, x, use_dropout=False):
        if use_dropout:
            x = torch.dropout(x, 0.5, train=False)
        return x

model = MyModel()
example_input = torch.randn(1, 3, 224, 224)

# 注意：这里 use_dropout 是 False
traced_model = torch.jit.trace(model, example_input)

# 问题来了：如果调用时 use_dropout=True
# traced_model 不会执行 dropout 分支！

我在项目中遇到过类似的问题。有一次导出一个带条件分支的模型，trace 出来的结果在特定输入下直接崩了。排查了半天，才发现是 trace 没记录到那个分支。

那 torch.jit.script 呢？它会分析你的 Python 代码，把整个控制流都编译进去。代价是——有些 Python 语法它不支持，比如动态的字典操作。

# script 方式
@torch.jit.script
def my_script_fn(x: torch.Tensor, n: int) -> torch.Tensor:
    result = torch.zeros_like(x)
    for i in range(n):
        result += x * i
    return result

# 这个就能正确处理循环了
scripted_fn = my_script_fn

我的建议：能用 trace 就用 trace，它简单可靠。遇到控制流复杂的模型，再考虑 script。如果 script 报语法错误，可以尝试把复杂逻辑拆成子模块。

torch.onnx.export：跨框架的桥梁

ONNX 是模型部署的通用语言。PyTorch 导出 ONNX 的核心函数就是 torch.onnx.export。

import torch.onnx

model = torchvision.models.resnet18(pretrained=True)
dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,               # 模型
    dummy_input,         # 示例输入
    "resnet18.onnx",     # 输出路径
    export_params=True,  # 导出参数
    opset_version=11,    # ONNX opset 版本
    do_constant_folding=True,  # 常量折叠优化
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

这里有个关键参数——opset_version。我刚开始做 ONNX 导出时，习惯用最新的 opset。后来发现，目标推理引擎可能不支持那么新的算子。比如 TensorRT 8.x 对 opset 13 支持得最好，opset 17 就有不少算子不兼容。

我曾经踩过的坑：导出 ONNX 时用了 opset 17，结果在 TensorRT 里报算子不支持。后来降到 opset 13，一切正常。建议先查目标引擎支持的 opset 版本，再决定导出参数。

动态轴（dynamic_axes）也很重要。如果你的模型需要处理不同 batch size 的输入，一定要设置这个。否则导出的 ONNX 会固定 batch size 为 1。

TensorFlow 模型导出：三种主流格式

TensorFlow 的导出方式比 PyTorch 复杂一些。我把它归纳为三种：SavedModel、Frozen Graph、TF-TRT。

SavedModel：TF 的标准格式

SavedModel 是 TensorFlow 官方推荐的导出格式。它包含了模型的计算图、权重、以及 serving 相关的签名信息。

import tensorflow as tf

# 构建模型
model = tf.keras.Sequential([
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 训练...（省略）

# 导出 SavedModel
model.save('saved_model/my_model', save_format='tf')

# 或者用 tf.saved_model.save
tf.saved_model.save(model, 'saved_model/my_model_v2')

SavedModel 的好处是完整——它把模型的训练状态、优化器参数都保存了。但部署时我们通常只需要推理，所以会去掉训练相关的部分。

实际经验：SavedModel 在 TF Serving 里用得最多。如果你用 TensorFlow Serving 部署，直接导出 SavedModel 就行，省事。

Frozen Graph：轻量级部署格式

Frozen Graph 是把计算图和权重合并成一个文件。它去掉了训练相关的节点，体积更小，加载更快。

import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

# 加载 SavedModel
loaded = tf.saved_model.load('saved_model/my_model')
infer = loaded.signatures['serving_default']

# 转换为 frozen graph
frozen_func = convert_variables_to_constants_v2(infer)
graph_def = frozen_func.graph.as_graph_def()

# 保存
with tf.io.gfile.GFile('model_frozen.pb', 'wb') as f:
    f.write(graph_def.SerializeToString())

我记得有一次在嵌入式设备上部署模型，SavedModel 加载要 3 秒，换成 Frozen Graph 后只要 0.5 秒。原因很简单——省去了解析训练节点的开销。

适用场景：Frozen Graph 特别适合资源受限的环境，比如手机端、边缘设备。它也是 TensorFlow Lite 转换的常用中间格式。

TF-TRT：NVIDIA 的加速方案

TF-TRT 是 TensorFlow 与 TensorRT 的集成方案。它可以在 TensorFlow 内部调用 TensorRT 的优化器，对模型进行层融合、精度校准等操作。

import tensorflow as tf
from tensorflow.python.compiler.tensorrt import trt_convert as trt

# 配置 TF-TRT 参数
converter = trt.TrtGraphConverterV2(
    input_saved_model_dir='saved_model/my_model',
    precision_mode='FP16',        # FP16 精度
    max_workspace_size_bytes=1 << 30,  # 1GB 工作空间
    maximum_cached_engines=100
)

# 转换
converter.convert()

# 保存
converter.save('trt_model')

TF-TRT 的好处是——你不需要把模型导出成 ONNX 再转 TensorRT。直接在 TensorFlow 里就能享受 TensorRT 的加速。但代价是灵活性受限，有些自定义算子可能不支持。

注意：TF-TRT 的精度模式选择要谨慎。FP16 能带来 2 倍左右的加速，但某些模型精度会下降。我建议先用 FP32 跑一遍 baseline，再对比 FP16 的结果。

三种导出方式的对比

格式 适用场景 优点 缺点
SavedModel TF Serving、训练后继续训练 完整保存所有信息 体积大、加载慢
Frozen Graph 移动端、边缘设备、TFLite 体积小、加载快 无法继续训练
TF-TRT NVIDIA GPU 推理加速 自动优化、集成方便 依赖 GPU、算子受限

知识体系总览

下面这张图展示了模型导出的整体流程和各个格式之间的关系：

模型导出工具链全景 PyTorch 训练 TensorFlow 训练 PyTorch 导出方式 • torch.jit.trace / script • torch.onnx.export → TorchScript / ONNX TensorFlow 导出方式 • SavedModel • Frozen Graph • TF-TRT 中间表示层 TensorRT / ONNX Runtime TFLite / Core ML TF Serving / Triton 图：模型导出流程 — 从训练框架到部署环境的中间表示转换

从这张图可以看得很清楚：无论你用 PyTorch 还是 TensorFlow，最终都要经过中间表示层，才能到达目标部署环境。我个人建议，团队里最好统一一种中间格式，比如 ONNX，这样不管前端用什么框架训练，后端都能对接。

好了，模型导出这块就聊到这儿。记住一点：导出不是终点，而是部署的起点。选对导出方式，后面能省很多事。

公众号:蓝海资料掘金营，微信deep3321

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