4. 模型序列化与格式：SavedModel、ONNX、TorchScript、TensorRT Plan 文件

模型序列化是将训练好的模型权重、计算图与元数据打包为单一文件或目录的过程。在低延迟因子评分场景中，序列化格式直接决定了推理引擎的兼容性、执行效率与部署灵活性。本章深入分析四种主流格式——SavedModel、ONNX、TorchScript 与 TensorRT Plan 文件——的结构特性、性能差异与最佳适用场景。

核心原则：序列化格式的选择应遵循“训练框架 → 中间表示 → 目标推理引擎”的转换链路。不要将训练格式直接用于生产推理，除非推理引擎原生支持且无性能损失。

4.1 格式总览与核心差异

下表从序列化结构、运行时依赖、算子覆盖与优化程度四个维度对比四种格式：

维度 SavedModel (TF) ONNX TorchScript TensorRT Plan
序列化结构 目录（pb + variables + assets） 单文件 .onnx（protobuf） 单文件 .pt（zip 归档） 单文件 .plan（二进制）
运行时依赖 TensorFlow Serving / TF Runtime ONNX Runtime / 各框架适配器 LibTorch / PyTorch JIT TensorRT / Triton 插件
算子覆盖 TF 全算子 跨框架标准算子集 PyTorch ATen 子集 TensorRT 优化算子（受限）
优化程度 图优化（Grappler） 图优化 + 量化（Q/DQ） JIT 融合 + 算子折叠 内核自动调优 + 层融合 + FP16/INT8
推理延迟（典型） 基线（1x） 0.8x – 1.2x 0.7x – 1.0x 0.3x – 0.6x
延迟敏感场景选型提示：若延迟要求 < 5ms，优先考虑 TensorRT Plan；若需跨框架部署且延迟容忍度在 10ms 以内，ONNX 是最佳中间格式；若团队统一使用 PyTorch 且无跨平台需求，TorchScript 可减少转换开销。

4.2 SavedModel：TensorFlow 生态的完整交付单元

SavedModel 是 TensorFlow 官方推荐的序列化格式，包含完整的计算图（SavedModel.pb）、变量检查点（variables/）与资产文件（assets/）。其核心优势在于版本管理与服务端签名——每个模型可附带多个 SignatureDef，明确指定输入/输出的张量名称与数据类型。

# 导出 SavedModel 示例
import tensorflow as tf

model = tf.keras.Sequential([...])
model.compile(...)
model.fit(x_train, y_train, epochs=5)

# 导出时指定签名
tf.saved_model.save(
    model,
    export_dir="./factor_model/1",
    signatures={
        'serving_default': model.call.get_concrete_function(
            tf.TensorSpec(shape=[None, 128], dtype=tf.float32, name="input_features")
        )
    }
)

优点：

  • 原生支持 TensorFlow Serving 的模型热加载与版本路由
  • 保留完整的训练元数据（优化器状态、损失函数配置）
  • 支持分布式部署中的模型分片

缺点：

  • 文件体积大（包含变量与资产）
  • 跨框架部署需额外转换
  • 推理延迟受限于 TF Runtime 的执行引擎

适用场景：纯 TensorFlow 技术栈、使用 TensorFlow Serving 做模型管理、需要版本回滚与 A/B 测试的在线评分系统。

4.3 ONNX：跨框架互操作的中间表示

ONNX（Open Neural Network Exchange）定义了一组标准算子与计算图格式，旨在打破框架壁垒。其核心价值在于一次导出，多处运行——从 PyTorch、TensorFlow、Keras 等框架导出的 ONNX 模型，可在 ONNX Runtime、TensorRT、OpenVINO 等推理后端上执行。

# PyTorch → ONNX 导出示例
import torch
import torch.onnx

model = MyFactorModel()
model.eval()
dummy_input = torch.randn(1, 128)

torch.onnx.export(
    model,
    dummy_input,
    "factor_model.onnx",
    input_names=["input_features"],
    output_names=["score"],
    dynamic_axes={"input_features": {0: "batch_size"}},
    opset_version=17
)

优点：

  • 跨框架、跨硬件的最佳互操作格式
  • 支持量化感知训练（QAT）与动态量化
  • 丰富的图优化工具链（onnx-simplifier、onnxoptimizer）

缺点：

  • 算子覆盖存在边界（部分自定义算子需实现适配器）
  • 动态形状支持有限（需显式声明 dynamic_axes）
  • 转换过程可能引入精度损失（尤其是控制流算子）

适用场景：多框架混合部署、需要将模型迁移至不同推理引擎、边缘端与云端统一模型格式。

ONNX 转换避坑：当模型中包含 if-else 控制流、循环或动态张量形状时，ONNX 导出可能失败或产生不可执行的图。建议在导出前使用 torch.jit.trace 或 tf.function 将动态图静态化。

4.4 TorchScript：PyTorch 的原生生产格式

TorchScript 是 PyTorch 的中间表示，通过 torch.jit.trace（跟踪）或 torch.jit.script（脚本化）将动态图转换为静态图。其核心优势在于无 Python 依赖——序列化后的 .pt 文件可在 C++ 运行时（LibTorch）中直接加载推理，彻底消除 GIL 与 Python 解释器开销。

# TorchScript 导出（跟踪模式）
import torch

model = MyFactorModel().eval()
example_input = torch.randn(1, 128)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("factor_model_traced.pt")

# 脚本化模式（支持控制流）
@torch.jit.script
def score_fn(features: torch.Tensor) -> torch.Tensor:
    if features.shape[0] > 100:
        return model(features) * 1.1
    return model(features)

torch.jit.save(score_fn, "factor_model_scripted.pt")

优点：

  • 与 PyTorch 训练流程无缝衔接，无需额外转换
  • 支持 C++ 部署，延迟极低（无 Python 开销）
  • JIT 编译器可自动执行算子融合与常量折叠

缺点：

  • 仅限 PyTorch 生态，跨框架需转 ONNX
  • 脚本化模式对 Python 动态特性支持有限
  • 跟踪模式可能丢失动态控制流

适用场景：纯 PyTorch 技术栈、需要 C++ 级低延迟推理、模型结构相对固定且无复杂控制流。

4.5 TensorRT Plan：GPU 推理的极致优化

TensorRT Plan 文件是 NVIDIA TensorRT 的专有二进制格式，包含经过层融合、精度校准与内核自动调优后的推理引擎。其优化过程将计算图转换为 GPU 上最高效的执行序列，支持 FP16、INT8 甚至 INT4 量化，是低延迟因子评分场景的终极选择。

# ONNX → TensorRT Plan 构建示例（Python API）
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)

with open("factor_model.onnx", "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
config.set_flag(trt.BuilderFlag.FP16)  # 启用 FP16

plan = builder.build_serialized_network(network, config)
with open("factor_model.plan", "wb") as f:
    f.write(plan)

优点：

  • GPU 推理延迟最低（相比 ONNX Runtime 可降低 40%-70%）
  • 支持 INT8 量化，吞吐量可提升 2-4 倍
  • 动态形状与多流执行优化

缺点：

  • 构建时间长（需 GPU 设备，可能耗时数分钟至数小时）
  • 算子支持受限（不支持的算子需自定义插件）
  • 硬件绑定（不同 GPU 架构需重新构建 Plan）

适用场景：高吞吐低延迟的 GPU 在线评分、批量推理延迟敏感型应用、已确定 GPU 型号的生产环境。

4.6 格式选择决策流程

下图展示了从训练框架到生产推理的格式选择路径：

模型序列化格式选择决策流程 训练框架 单一框架？跨框架？ TF → SavedModel PyTorch → TorchScript 导出 ONNX GPU 推理？延迟要求？ TensorRT Plan ONNX Runtime 生产推理服务

4.7 低延迟场景的格式组合策略

在实际的因子评分系统中，单一格式往往无法满足所有需求。推荐采用多格式共存策略：

  1. 训练阶段：使用框架原生格式（SavedModel / .pt）保存完整训练状态
  2. 验证阶段：导出 ONNX 进行跨框架精度对齐测试
  3. 生产阶段：根据部署硬件选择 TensorRT Plan（GPU）或 ONNX Runtime（CPU）
  4. 回退方案：保留 TorchScript 或 SavedModel 作为降级推理路径
延迟预算分配建议：若端到端延迟预算为 10ms，建议模型推理占 6-7ms，序列化/反序列化占 1-2ms，网络传输占 1-2ms。选择 TensorRT Plan 可将推理部分压缩至 2-3ms，为其他环节留出余量。

4.8 格式转换的精度验证

每次格式转换都必须进行数值精度验证。推荐使用余弦相似度与最大相对误差两个指标：

import numpy as np

def validate_conversion(original_model, converted_model, test_inputs):
    """验证转换前后模型输出的精度损失"""
    orig_outputs = original_model(test_inputs).detach().numpy()
    conv_outputs = converted_model(test_inputs).detach().numpy()

    # 余弦相似度
    cos_sim = np.dot(orig_outputs.flatten(), conv_outputs.flatten()) / (
        np.linalg.norm(orig_outputs) * np.linalg.norm(conv_outputs)
    )

    # 最大相对误差
    max_rel_error = np.max(np.abs(orig_outputs - conv_outputs) / 
                          (np.abs(orig_outputs) + 1e-8))

    print(f"余弦相似度: {cos_sim:.6f}")
    print(f"最大相对误差: {max_rel_error:.6f}")

    # 因子评分场景通常要求余弦相似度 > 0.999，最大相对误差 < 0.01
    assert cos_sim > 0.999, "精度损失超出容忍范围"
    assert max_rel_error < 0.01, "相对误差过大"
量化精度陷阱：INT8 量化在因子评分模型中可能导致尾部风险预测偏差。建议在量化后使用 KS 检验或分位数对比验证评分分布的一致性，而不仅仅是平均精度指标。
← 上一章 📖 返回目录 下一章 →