问题背景
先丢一组我自己踩过的“血淋淋”数据：
- 训练 3 天的 CosyVoice 多说话人模型，因为
torch.save时把model.state_dict和optimizer.state_dict混在一个文件里，结果线上推理加载失败，直接回滚，浪费 120 张 A100 卡时 ≈ 1.4 w 元； - 某次热更新，忘了把
model.eval()写进保存脚本，BatchNorm 统计量带噪，WER 从 6.8% 飙到 11.2%，客诉率 +300%； - 把 1.1 GB 的 float32 模型直接丢进移动端，加载耗时 18 s，内存峰值 2.3 GB，直接被系统守护进程 kill。
一句话：模型保存不是“ctrl+s”那么简单，格式、精度、版本、存储、安全，每一步都能埋雷。
方案选型
我把团队过去一年试过的三种主流路线拉出来对比，结论先看表，后面再给代码。
| 方案 | 适用场景 | 关键参数 | 体积(相对) | 首次推理延迟 | 跨平台 | 备注 |
|---|---|---|---|---|---|---|
| torch.save(state_dict) | 继续训练/热启动 | pickle_protocol=4,_use_new_zipfile=True | 1× | 高 | 仅限 Python | 最灵活，也最容易埋雷 |
| ONNX (opset=14+) | 服务化 CPU/GPU 推理 | do_constant_folding=True,export_params=True | 0.7× | 中 | 需固定输入 shape，动态轴要手动声明 | |
| TensorRT (fp16/int8) | 边缘端实时推理 | max_workspace_size=1<<30,fp16=True | 0.35× | 最低 | 构建耗时，对 CUDA 版本敏感 |
经验：训练阶段用
torch.save做 checkpoint，上线前转 ONNX，边缘盒再转 TensorRT，基本“一鱼三吃”。
，下面给出可抄作业的脚本。
核心实现
以下代码全部 PEP8，双语注释，复制即可跑。示例以 CosyVoice 的CosyVoiceASR类为原型，其他结构只需把model换成自己的实例。
1. 通用保存函数（含校验）
# save_model.py import os import hashlib import torch import onnx from typing import Dict, Optional def save_checkpoint( model: torch.nn.Module, optimizer: torch.optim.Optimizer, epoch: int, ckpt_dir: str = "./ckpt", max_keep: int = 5, ) -> str: """ 保存训练状态，返回文件路径 Save training state, return file path """ os.makedirs(ckpt_dir, exist_ok=True) ckpt_path = os.path.join(ckpt_dir, f"model_ep{epoch:04d}.pth") state = { "epoch": epoch, "model_state": model.state_dict(), "optim_state": optimizer.state_dict(), } torch.save(state, ckpt_path, _use_new_zipfile_serialization=True) _remove_old_ckpt(ckpt_dir, max_keep) return ckpt_path def save_inference_model( model: torch.nn.Module, dummy_input: torch.Tensor, save_path: str, export_onnx: bool = True, opset: int = 14, ) -> None: """ 导出推理用模型，默认同时保存 .pth 和 .onnx Export inference model, default save both .pth & .onnx """ model.eval() # 必须，否则 BN 会带噪 with torch.no_grad(): # --- 1. torch.save 仅保留权重 --- torch.save(model.state_dict(), save_path + ".pth") # --- 2. ONNX 导出 --- if export_onnx: onnx_path = save_path + ".onnx" torch.onnx.export( model, dummy_input, onnx_path, input_names=["speech"], output_names=["text_logits"], dynamic_axes={"speech": {0: "batch", 1: "seq"}, "text_logits": {0: "batch", 1: "seq"}}, opset_version=opset, do_constant_folding=True, export_params=True, ) # 简单校验 onnx_model = onnx.load(onnx_path) onnx.checker.check_model(onnx_model) print(f"[INFO] ONNX exported & checked: {onnx_path}")2. 加载 + 一致性校验
# load_model.py import torch import onnxruntime as ort import numpy as np from typing import Dict def load_torch_model( model_kls: torch.nn.Module, weight_path: str, device: str = "cpu", ) -> torch.nn.Module: """ 加载权重并校验键值匹配 Load weights and check key matching """ model = model_kls() state = torch.load(weight_path, map_location=device) model.load_state_dict(state, strict=True) model.eval() return model def load_onnx_model(onnx_path: str) -> ort.InferenceSession: """ 加载 ONNX 并返回 runtime session """ sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) return sess def consistency_check( torch_model: torch.nn.Module, onnx_sess: ort.InferenceSession, dummy: np.ndarray, rtol: float = 1e-3, ) -> bool: """ 比较 torch & onnx 输出误差 Compare torch & onnx output error """ torch_out = torch_model(torch.from_numpy(dummy)).detach().numpy() onnx_out = onnx_sess.run(None, {"speech": dummy})[0] flag = np.allclose(torch_out, onnx_out, rtol=rtol) print(f"[INFO] consistency check {'passed' if flag else 'failed'}") return flag3. 量化压缩（TensorRT 示例）
# trt_convert.py import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit # 初始化 CUDA context def build_engine(onnx_path: str, max_batch=8, fp16=True) -> trt.ICudaEngine: """ 将 ONNX 转为 TensorRT engine，返回序列化引擎 """ logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) config = builder.create_builder_config() config.max_workspace_size = 1 << 30 # 1 GB if fp16: config.set_flag(trt.BuilderFlag.FP16) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) parser = trt.OnnxParser(network, logger) with open(onnx_path, "rb") as f: parser.parse(f.read()) engine = builder.build_engine(network, config) return engine def save_engine(engine: trt.ICudaEngine, path: str): with open(path, "wb") as f: f.write(engine.serialize())注意：TensorRT 构建过程慢，建议在 CI 里预编译并缓存
.engine文件，别在线转。
生产考量
版本控制
用model-epoch-xxx.sha256的方式，把权重文件、onnx、engine 和一份meta.yaml（记录训练 git commit、数据集 md5、wer）一起打 tar，上传至 S3/OSS。
文件名带epoch + git short sha，再也不怕“谁动了我的模型”。加密存储
边缘盒子怕被抄？把.engine用 AES-CTR 加密，启动时通过 TPM 解密到内存，不落盘。Python 端可直接调pycryptodome，解密后走trt.Runtime反序列化。边缘部署
板子内存只有 4 GB？先fp16，再开layerwise fusion，最后把max_workspace_size压到 256 MB。实测 CosyVoice 从 1.1 GB 降到 380 MB，首帧延迟 120 ms → 45 ms。热更新
起双进程 + 共享内存队列，A 进程加载旧模型，B 进程加载新模型，B 初始化完通过 unix socket 告诉 A 切换指针，实现 0 downtime。记得加读写锁，防止并发时序错乱。
避坑指南
未冻结 BN / Dropout
保存前务必model.eval()，否则统计量会跟着 batch 跑，推理结果随机“漂移”。动态轴忘声明
ONNX 导出时dynamic_axes不填，默认 batch=1，线上多路并发直接报错。把 optimizer 一起打包
checkpoint 里留optim_state没问题，但上线推理模型别带它，体积翻倍，还易泄露学习率等敏感信息。TensorRT 版本不一致
构建和运行时的 CUDA/cuDNN/TRT 版本要锁死，差一个小版本都可能deserialize失败。忘记一致性校验
转完 ONNX/TensorRT 一定跑np.allclose，误差>1% 就打回重训，别等线上用户帮你发现“词错率”翻倍。
写在最后
模型保存听起来像收尾工作，实则是一条“最后一公里”的长链：格式、精度、体积、安全、版本、灰度，每一步都能让前面的训练成本打水漂。上面这套脚本和 checklist 已经在我们内部流水线跑了 8 个月，累计 200+ 次安全上线，存储开销平均降 40%，回滚次数从每月 3 次降到 0 次。
但问题还没完：当实验规模膨胀到“日更 50 个模型”时，人工打 tar、上传、写 YAML 迟早会失控。如何设计一套自动化模型归档系统，让 CI 自动完成“训练→评估→对比基线→生成指纹→加密→推送边缘”？欢迎评论区一起头脑风暴，也许下一个开源项目就诞生在你的回复里。
