4. 模型部署格式：ONNX 导出、TorchScript、TensorRT Engine 的生成与验证

模型训练好只是第一步，真正让它跑起来，还得过部署这一关。说白了，就是把 PyTorch 或者 TensorFlow 的模型，转成推理引擎能吃的格式。我这些年踩过的坑，有一半都跟格式转换有关。今天咱们就把 ONNX、TorchScript、TensorRT Engine 这三个主流格式聊透。

4.1 为什么需要中间表示？

你想想看，训练框架追求的是灵活，能动态建图、能自动求导。但推理引擎要的是极致性能——算子融合、内存复用、量化压缩。这两者天然矛盾。

所以我们需要一个“中间人”。它把训练好的模型冻结成静态图，去掉反向传播的包袱，只保留前向推理的骨架。这个中间人，就是部署格式。

核心原则：训练框架负责“能跑”，部署格式负责“跑得快”。两者解耦，各司其职。

4.2 ONNX 导出：通用交换格式

ONNX（Open Neural Network Exchange）是微软和 Meta 搞的开放标准。它的优势在于生态广——几乎所有推理引擎都支持导入 ONNX。我个人习惯把 ONNX 当作“第一站”，先把模型转成 ONNX，再分发到不同后端。

4.2.1 PyTorch 导出 ONNX

PyTorch 提供了 torch.onnx.export 接口。看起来简单，但参数细节很多。我直接给一个生产可用的模板：

import torch
import torch.onnx

def export_to_onnx(model, dummy_input, output_path):
    model.eval()
    # 设置动态轴：batch size 可变
    dynamic_axes = {
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
    
    torch.onnx.export(
        model,
        dummy_input,
        output_path,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes=dynamic_axes,
        opset_version=17,          # 推荐 17 以上
        do_constant_folding=True,  # 常量折叠优化
        export_params=True,        # 导出权重
        verbose=False
    )
    print(f"ONNX 模型已保存至: {output_path}")

小技巧：opset_version 别用太老。我建议至少 17，因为新版本支持更多算子，比如 Transformer 里的 LayerNorm 和 Softmax 都能原生导出，不用走 fallback 路径。

4.2.2 验证 ONNX 模型

导出完一定要验证。我曾经吃过一次亏——模型导出时静默失败，算子被替换成了不支持的版本，推理结果全错。所以验证步骤不能省。

import onnx
import onnxruntime as ort
import numpy as np

# 1. 检查模型结构
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print("ONNX 结构检查通过")

# 2. 运行推理对比精度
ort_session = ort.InferenceSession("model.onnx")
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = ort_session.run(None, {'input': input_data})

# 3. 与原始 PyTorch 结果对比
with torch.no_grad():
    torch_outputs = model(torch.from_numpy(input_data))
    
diff = np.abs(outputs[0] - torch_outputs.numpy()).max()
print(f"最大误差: {diff:.6f}")
assert diff < 1e-4, "精度偏差过大！"

注意：ONNX 验证时，输入数据的 dtype 和 shape 必须与导出时完全一致。尤其是 batch size 为 1 时，别忘了保持维度是 4D 张量。

4.3 TorchScript：PyTorch 原生部署方案

TorchScript 是 PyTorch 自家的部署格式。它有两种生成方式：Tracing 和 Scripting。我一般混合使用——能用 Tracing 的用 Tracing，遇到控制流再用 Scripting。

4.3.1 Tracing 方式

Tracing 就是给模型一个 dummy input，跟踪执行路径，记录下计算图。优点是简单，缺点是有动态分支时容易漏掉路径。

import torch

class MyModel(torch.nn.Module):
    def forward(self, x, use_dropout=False):
        x = self.conv(x)
        if use_dropout:  # 这个分支 Tracing 会漏掉
            x = self.dropout(x)
        return x

model = MyModel()
dummy = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, dummy)
traced_model.save("model_traced.pt")

4.3.2 Scripting 方式

Scripting 会解析 Python 源码，把控制流也编译进去。适合有 if/for 的模型。

scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

我的经验：对于 Transformer 这类模型，我建议用 Scripting。因为 self-attention 里有 mask 判断，Tracing 只能看到一条路径，推理时换一个 mask 长度就可能崩。

4.3.3 TorchScript 验证

验证 TorchScript 模型，核心是保证输出一致：

# 加载 TorchScript 模型
jit_model = torch.jit.load("model_scripted.pt")
jit_model.eval()

# 对比输出
with torch.no_grad():
    jit_out = jit_model(dummy)
    orig_out = model(dummy)
    
print(torch.allclose(jit_out, orig_out, atol=1e-5))

4.4 TensorRT Engine：NVIDIA 的终极加速

TensorRT 是 NVIDIA 的推理优化引擎。它能把 ONNX 模型编译成高度优化的 Engine，支持 FP16、INT8 量化，还能做层融合和内存优化。说白了，这是把 GPU 性能榨干的最后一步。

4.4.1 ONNX 转 TensorRT Engine

我推荐用 trtexec 命令行工具，简单直接：

# 基本用法：ONNX 转 FP16 Engine
trtexec --onnx=model.onnx \
        --saveEngine=model_fp16.engine \
        --fp16 \
        --workspace=4096 \
        --minShapes=input:1x3x224x224 \
        --optShapes=input:8x3x224x224 \
        --maxShapes=input:32x3x224x224

参数说明：

参数 作用 建议值
--fp16 启用半精度推理 V100/T4/A100 必开
--int8 INT8 量化（需校准数据） 精度要求不高时使用
--workspace 编译时可用显存 (MB) 4096 ~ 8192
--minShapes/optShapes/maxShapes 动态 batch 范围 根据业务流量定

避坑：动态 shape 的 min/opt/max 必须设置合理。我见过有人把 min 设成 1，max 设成 128，结果 Engine 编译了 40 分钟。因为 TensorRT 要为每个可能的 shape 都做优化。建议 min 和 max 差距不要超过 8 倍。

4.4.2 Python API 方式生成

如果你需要在代码里动态生成 Engine，可以用 Python API：

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_path, engine_path, fp16=True):
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network(
        1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
    )
    parser = trt.OnnxParser(network, TRT_LOGGER)
    
    with open(onnx_path, 'rb') as f:
        parser.parse(f.read())
    
    config = builder.create_builder_config()
    config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 << 30)  # 4GB
    
    if fp16 and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    
    # 设置动态 shape
    profile = builder.create_optimization_profile()
    profile.set_shape('input', (1,3,224,224), (8,3,224,224), (32,3,224,224))
    config.add_optimization_profile(profile)
    
    engine = builder.build_serialized_network(network, config)
    with open(engine_path, 'wb') as f:
        f.write(engine)
    print(f"Engine 已保存至: {engine_path}")

4.4.3 TensorRT Engine 验证

验证 Engine 时，我习惯做三件事：结构检查、精度对比、性能测试。

import pycuda.driver as cuda
import pycuda.autoinit

def validate_engine(engine_path, dummy_input):
    # 加载 Engine
    with open(engine_path, 'rb') as f:
        runtime = trt.Runtime(TRT_LOGGER)
        engine = runtime.deserialize_cuda_engine(f.read())
    
    # 创建执行上下文
    context = engine.create_execution_context()
    
    # 分配 GPU 内存
    input_idx = engine.get_binding_index('input')
    output_idx = engine.get_binding_index('output')
    
    # 设置动态 shape（如果用了动态 batch）
    context.set_binding_shape(input_idx, dummy_input.shape)
    
    # 执行推理
    output = np.empty_like(dummy_input)
    # ... 省略内存拷贝和 kernel launch 代码
    
    # 精度对比
    with torch.no_grad():
        torch_out = model(torch.from_numpy(dummy_input))
    
    mse = np.mean((output - torch_out.numpy())**2)
    print(f"TensorRT vs PyTorch MSE: {mse:.6f}")
    assert mse < 1e-3, "精度偏差过大！"

4.5 三种格式对比

特性 ONNX TorchScript TensorRT Engine
生态兼容性 最广（几乎所有框架支持） 仅 PyTorch 生态 仅 NVIDIA GPU
推理速度 中等（依赖运行时） 中等（PyTorch JIT） 最快（硬件级优化）
量化支持 需额外工具 有限 原生 FP16/INT8
动态 shape 支持（需设置 dynamic_axes） 支持 支持（需 profile）
部署复杂度 低 低 高（需 GPU 环境）

4.6 部署格式转换流程

下面这张图是我在实际项目中总结的转换链路，你可以把它当作标准流程来用。

模型部署格式转换流程 PyTorch 训练模型 torch.onnx.export ONNX 模型 trtexec torch.jit.trace/script TensorRT Engine TorchScript 模型 生产环境部署（Triton / TorchServe / 自研引擎） 💡 推荐路径：PyTorch → ONNX → TensorRT Engine（追求极致性能） 或 PyTorch → TorchScript（快速部署，不依赖 GPU 环境）

4.7 常见问题与避坑

  • 算子不支持：ONNX 导出时遇到不支持的算子，可以用 torch.onnx.register_custom_op 注册自定义实现。我遇到过 GroupNorm 导出失败，最后自己写了个 ONNX 算子才搞定。
  • 动态 shape 性能差：TensorRT 的动态 shape 会引入额外开销。如果业务流量稳定，建议用固定 shape 编译 Engine，性能能提升 10%~20%。
  • 精度下降：FP16 推理时，某些层（比如 LayerNorm）的精度敏感。我建议对精度敏感层保持 FP32，用 TensorRT 的 set_precision 接口做混合精度。
  • 内存泄漏：TensorRT Engine 反复创建/销毁时，CUDA context 可能泄漏。我习惯用对象池复用 Engine，避免频繁创建。

我的建议：别一上来就追求极致优化。先用 ONNX 跑通流程，再用 TensorRT 做性能压榨。每一步都做验证，出了问题能快速定位。我曾经跳过 ONNX 验证直接转 TensorRT，结果花了三天才找到是某个算子精度不对——从那以后，我再也不敢省验证步骤了。

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