📚 课程目录第3章模型导出与转换

3. 模型导出与转换

模型训练好只是第一步，真正要让它跑在推理服务里，还得过一道关——模型导出与转换。说白了，就是把PyTorch那套动态图，变成生产环境能高效执行的静态图。

我个人习惯把这一步叫做「模型工业化」。你想想看，训练时我们用PyTorch，灵活、好调试，但到了线上，我们需要的是速度、是稳定。所以就有了TorchScript、ONNX、TensorRT这些中间格式。

核心思路：训练框架（PyTorch）→ 中间表示（TorchScript/ONNX）→ 推理引擎（TensorRT）。每一步都有坑，但每一步也都有收益。

TorchScript：PyTorch的原生序列化方案

TorchScript是PyTorch官方推出的模型序列化格式。它能把动态图「凝固」成静态图，脱离Python运行时也能跑。

我记得第一次用TorchScript是在一个文本分类项目里。模型在开发机上跑得好好的，一上生产环境就报错——因为生产环境没装PyTorch的完整版。后来改用TorchScript导出，问题就解决了。

TorchScript有两种导出方式：

方式	适用场景	优点	缺点
跟踪（Trace）	模型不含控制流（if/for依赖数据）	简单、快速、一次运行即可导出	无法处理动态分支
脚本（Script）	模型含控制流或动态结构	完整保留控制流逻辑	需要写类型注解，部分操作不支持

来看代码。先说说Trace方式：

import torch
import torchvision.models as models

# 加载预训练模型
model = models.resnet18(pretrained=True)
model.eval()

# 构造示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 使用 torch.jit.trace 导出
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save("resnet18_traced.pt")

# 加载验证
loaded = torch.jit.load("resnet18_traced.pt")
output = loaded(dummy_input)
print(output.shape)  # torch.Size([1, 1000])

嗯，这里要注意。Trace方式有个大坑——它只记录实际执行过的路径。如果你的模型里有if x.shape[0] > 10这样的分支，Trace只会记录其中一条分支。我曾经因为这个原因，线上模型在特定输入下直接崩溃。

那Script方式怎么用？

import torch

class MyModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(10, 5)
    
    def forward(self, x):
        # 包含控制流
        if x.sum() > 0:
            return self.linear(x)
        else:
            return torch.zeros(5)

model = MyModel()
model.eval()

# 使用 torch.jit.script 导出
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

我的建议：能Trace就别Script。Trace更稳定，兼容性更好。如果模型必须用Script，记得给所有张量操作加上类型注解，否则容易报错。

ONNX：跨框架的通用中间格式

ONNX（Open Neural Network Exchange）是微软和Meta联合推出的开放格式。它的目标很明确——让模型在不同框架之间自由迁移。

我个人觉得，ONNX最大的价值在于「解耦」。你想想看，训练用PyTorch，推理用TensorRT，中间靠什么衔接？就是ONNX。它像一个万能转接头。

导出ONNX的代码其实很简单：

import torch
import torchvision.models as models

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

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

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

print("ONNX导出完成")

这里有个关键参数——dynamic_axes。它允许输入输出的batch size是动态的。我在项目中遇到过，如果不设置这个，导出的ONNX只能处理固定batch size的输入，线上服务一旦遇到不同batch的请求就直接报错。

注意：ONNX导出时，opset_version不是越高越好。高版本算子虽然功能多，但某些推理引擎（比如TensorRT）可能不支持。我一般用opset 11或13，兼容性最好。

导出后怎么验证？用ONNX Runtime跑一下：

import onnx
import onnxruntime as ort
import numpy as np

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

# 用ONNX Runtime推理
session = ort.InferenceSession("resnet50.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
result = session.run([output_name], {input_name: test_input})
print(f"输出形状: {result[0].shape}")

TensorRT：NVIDIA的推理加速引擎

TensorRT是NVIDIA推出的深度学习推理优化库。它能把ONNX模型编译成高度优化的引擎文件，支持FP16和INT8量化，推理速度能提升2-5倍。

我记得第一次用TensorRT是在一个实时视频分析项目里。原来用PyTorch直接推理，1080p的视频只能跑15fps。换成TensorRT的FP16模式后，直接飙到45fps。嗯，那个感觉，真爽。

TensorRT转换的基本流程：

准备ONNX模型——确保导出时算子兼容
创建Builder——配置精度、工作空间等参数
解析ONNX——用ONNX Parser读取模型
构建引擎——编译优化，生成.engine文件
序列化保存——保存引擎文件供后续加载

来看代码实现：

import tensorrt as trt
import numpy as np

# TensorRT日志记录器
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

def build_engine(onnx_file_path, engine_file_path, fp16_mode=False):
    """将ONNX模型转换为TensorRT引擎"""
    
    with trt.Builder(TRT_LOGGER) as builder, \
         builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) as network, \
         trt.OnnxParser(network, TRT_LOGGER) as parser:
        
        # 配置Builder
        config = builder.create_builder_config()
        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
        
        if fp16_mode and builder.platform_has_fast_fp16:
            config.set_flag(trt.BuilderFlag.FP16)
            print("启用FP16模式")
        
        # 解析ONNX
        with open(onnx_file_path, "rb") as f:
            if not parser.parse(f.read()):
                for error in range(parser.num_errors):
                    print(parser.get_error(error))
                raise RuntimeError("ONNX解析失败")
        
        print(f"网络输入: {network.num_inputs}")
        print(f"网络输出: {network.num_outputs}")
        
        # 构建引擎
        serialized_engine = builder.build_serialized_network(network, config)
        
        # 保存引擎文件
        with open(engine_file_path, "wb") as f:
            f.write(serialized_engine)
        
        print(f"引擎已保存至: {engine_file_path}")
        return serialized_engine

# 调用
build_engine("resnet50.onnx", "resnet50_fp16.engine", fp16_mode=True)

经验之谈：TensorRT构建引擎时，workspace大小很关键。设太小了，复杂模型可能构建失败；设太大了，显存不够用。我一般设1GB起步，不够再加。

引擎构建好后，加载推理：

import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np

def load_engine(engine_file_path):
    """加载TensorRT引擎"""
    with open(engine_file_path, "rb") as f, \
         trt.Runtime(TRT_LOGGER) as runtime:
        return runtime.deserialize_cuda_engine(f.read())

def inference(engine, input_data):
    """执行推理"""
    with engine.create_execution_context() as context:
        # 分配GPU内存
        output_size = 1
        for binding in engine:
            if engine.binding_is_input(binding):
                input_idx = engine[binding]
                input_shape = context.get_binding_shape(input_idx)
                input_size = trt.volume(input_shape)
                # 分配输入输出缓冲区
                d_input = cuda.mem_alloc(input_size * np.dtype(np.float32).itemsize)
            else:
                output_idx = engine[binding]
                output_shape = context.get_binding_shape(output_idx)
                output_size = trt.volume(output_shape)
                d_output = cuda.mem_alloc(output_size * np.dtype(np.float32).itemsize)
        
        # 拷贝数据到GPU
        cuda.memcpy_htod(d_input, input_data.astype(np.float32).ravel())
        
        # 执行推理
        context.execute_v2([int(d_input), int(d_output)])
        
        # 拷贝结果回CPU
        output_data = np.empty(output_size, dtype=np.float32)
        cuda.memcpy_dtoh(output_data, d_output)
        
        return output_data

# 使用
engine = load_engine("resnet50_fp16.engine")
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)
result = inference(engine, test_input)
print(f"推理结果形状: {result.shape}")

三种格式的对比与选型建议

特性	TorchScript	ONNX	TensorRT
依赖	仅PyTorch	ONNX Runtime / 各框架	NVIDIA GPU + CUDA
推理速度	中等	中等	最快（FP16/INT8）
跨平台	仅PyTorch生态	广泛	仅NVIDIA GPU
量化支持	有限	一般	原生支持
调试难度	低	中等	高

我的选型建议：

纯PyTorch部署，不跨框架 → TorchScript，省事
需要跨框架或跨硬件 → ONNX，通用性最强
追求极致性能，且用NVIDIA GPU → TensorRT，值得折腾
生产环境常见组合：PyTorch → ONNX → TensorRT，兼顾通用性和性能

最后说一句，模型导出这件事，看起来简单，但坑不少。我建议你每导出一个格式，都用对应的Runtime跑一遍验证。别等到上线了才发现模型跑不起来——那种感觉，嗯，我不想再体验第二次。