4、模型加载与初始化：序列化格式、加载方式与预热策略

模型训练好了，怎么把它请进推理环境？这一步看似简单，但坑不少。我见过不少团队，模型在训练时跑得飞起，一到生产就崩，十有八九是加载和初始化环节出了问题。

说白了，模型加载就是「把硬盘上的二进制文件，还原成内存里的计算图」。但不同的序列化格式、不同的硬件拓扑，处理方式天差地别。今天我们就聊聊这三件事：格式选型、加载方式、预热策略。

4.1 模型序列化格式：三足鼎立

目前主流格式就三个：ONNX、TorchScript、SavedModel。我个人的选择原则很简单——看你的部署后端。

4.1.1 ONNX：通用但需小心算子兼容

ONNX 的好处是「一次导出，到处运行」。但我在项目中遇到过一个问题：PyTorch 里的 torch.topk 导出成 ONNX 后，在 TensorRT 上跑不起来。原因是 TensorRT 的 ONNX 解析器不支持动态 k 值。

避坑指南：导出前先用 onnxruntime 跑一遍推理，验证结果是否对齐。我曾经因为跳过这步，上线后才发现精度偏差，回滚了一整个版本。

# 导出 ONNX 时的常见做法
import torch
import torch.onnx

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

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

# 验证
import onnxruntime as ort
session = ort.InferenceSession("model.onnx")
output = session.run(None, {"input": dummy_input.numpy()})
print(output[0].shape)

4.1.2 TorchScript：PyTorch 部署的首选

如果你整个链路都是 PyTorch，TorchScript 是最省心的。它有两种导出方式：trace 和 script。

我个人习惯用 trace 做快速验证，用 script 处理带控制流的模型。你想想看，如果你的模型里有 if 语句，trace 是抓不到的，必须用 script。

# TorchScript 导出
model = MyModel().eval()

# 方式一：trace（适合无分支模型）
traced_model = torch.jit.trace(model, torch.randn(1, 3, 224, 224))
traced_model.save("model_traced.pt")

# 方式二：script（适合有控制流）
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

# C++ 加载
# torch::jit::load("model_scripted.pt");

4.1.3 SavedModel：TF 生态的完整方案

SavedModel 不只是模型权重，它还保存了输入输出的签名信息。这意味着你不需要在服务端写额外的预处理代码。TF Serving 可以直接根据签名自动匹配请求。

# TensorFlow SavedModel 导出
import tensorflow as tf

model = tf.keras.models.load_model("training_model.h5")
tf.saved_model.save(model, "saved_model/1/")

# 查看签名
loaded = tf.saved_model.load("saved_model/1/")
print(loaded.signatures)

4.2 模型加载方式：单卡 vs 多卡

加载方式取决于你的推理服务器有几张卡。嗯，这里要注意：训练时的多卡和推理时的多卡，逻辑完全不同。

4.2.1 单卡加载：最直接的方式

单卡没什么好说的，就是 torch.load 然后 model.to(device)。但有一个细节：先加载到 CPU，再转移到 GPU。直接 map_location='cuda:0' 虽然可以，但如果你后续要做模型预热，CPU 加载更灵活。

# 推荐的单卡加载方式
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = MyModel()
checkpoint = torch.load("model.pt", map_location="cpu")
model.load_state_dict(checkpoint)
model.to(device)
model.eval()

4.2.2 多卡加载：数据并行 vs 模型并行

数据并行是最常见的。每张卡上放一份完整的模型，输入数据按 batch 维度切分。我建议用 torch.nn.DataParallel 做快速原型，生产环境用 torch.distributed 或者 Ray Serve。

模型并行则适用于超大模型（比如 70B 参数）。每张卡只放一部分层，推理时数据要串行经过所有卡。延迟会高一些，但能跑单卡放不下的模型。

# 多卡数据并行加载
model = MyModel()
checkpoint = torch.load("model.pt", map_location="cpu")
model.load_state_dict(checkpoint)

# 方式一：DataParallel（简单，但主卡负载高）
model = torch.nn.DataParallel(model).cuda()

# 方式二：DistributedDataParallel（推荐生产）
# 需要先初始化进程组
# torch.distributed.init_process_group(backend="nccl")
# model = torch.nn.parallel.DistributedDataParallel(model.cuda())

4.3 模型预热（Warm-up）策略

模型刚加载完，第一次推理通常很慢。为什么？因为 GPU 的 CUDA 内核是懒加载的。第一次调用时，驱动要编译 kernel、分配显存、建立缓存。这个过程可能耗时几百毫秒甚至几秒。

预热，就是提前跑几次推理，让这些「一次性开销」在服务启动时完成，而不是在第一个用户请求时发生。

4.3.1 预热的核心逻辑

说白了，预热就是「用假数据跑几遍推理」。但要注意：预热的数据形状必须和线上一致。如果线上是 224x224，预热用 256x256，那预热就白做了——CUDA 会重新编译 kernel。

# 模型预热示例
def warm_up(model, device, input_shape=(1, 3, 224, 224), num_iterations=10):
    """
    预热模型，确保 CUDA kernel 被编译
    """
    model.eval()
    dummy_input = torch.randn(*input_shape).to(device)
    
    # 先跑一次，触发 kernel 编译
    with torch.no_grad():
        _ = model(dummy_input)
    
    # 再跑几次，让缓存稳定
    for _ in range(num_iterations - 1):
        with torch.no_grad():
            _ = model(dummy_input)
    
    # 同步 CUDA，确保所有操作完成
    torch.cuda.synchronize()
    print(f"Warm-up completed with {num_iterations} iterations.")

# 使用
device = torch.device("cuda:0")
model = load_model().to(device)
warm_up(model, device, input_shape=(1, 3, 224, 224), num_iterations=10)

4.3.2 预热的最佳实践

我建议预热时考虑以下几点：

• 预热次数：3-10 次就够了。太多浪费启动时间，太少可能没完全触发。
• 预热数据：用随机噪声即可，不需要真实数据。但形状必须匹配。
• 多 batch 预热：如果线上有不同 batch size，每个 size 都要预热一次。
• 多卡预热：每张卡都要单独预热，因为每张卡的 CUDA context 是独立的。

4.3.3 预热与模型加载的整合

在实际部署中，我会把预热和加载写成一个启动脚本。服务启动时自动完成，用户无感知。

# 完整的加载+预热流程
def load_and_warm_up(model_path, device, batch_size=1):
    # 1. 加载模型
    model = MyModel()
    checkpoint = torch.load(model_path, map_location="cpu")
    model.load_state_dict(checkpoint)
    model.to(device)
    model.eval()
    
    # 2. 预热
    dummy_input = torch.randn(batch_size, 3, 224, 224).to(device)
    for _ in range(5):
        with torch.no_grad():
            _ = model(dummy_input)
    torch.cuda.synchronize()
    
    print(f"Model loaded and warmed up on {device}")
    return model

# 启动时调用
model = load_and_warm_up("model.pt", "cuda:0", batch_size=8)

4.4 本章知识体系

下面这张图总结了模型加载与初始化的核心流程。你可以把它当作一个检查清单：

模型加载与初始化，看似是部署中最简单的环节，但细节决定成败。格式选错了，加载方式不对，预热没做好——任何一个环节出问题，都会直接影响线上延迟和吞吐量。

我个人习惯在每次部署前，都跑一遍完整的加载+预热流程，记录下首次推理延迟和稳定后的延迟。如果差距超过 5 倍，说明预热没到位，需要调整。