加载一个预训练模型（或者换成你自己的）

2025-11-28 489 AI链物

一键打包你的AI模型：从训练到导出的完整指南

搞出了一个能准确识别猫狗图片的模型，或者训练出一个能写诗的小助手，结果卡在最后一步——不知道怎么把它拿出来用？别急，这事儿就像打包行李，看似简单,但装不好就容易丢三落四。

模型导出到底在搞什么名堂？

说白了，就是把训练好的模型从训练环境里“搬”出来，变成能在其他地方使用的格式，想象一下，你在自家厨房研究出一道绝味辣椒炒肉，现在要把这道菜的“配方”和“火候秘诀”记录下来,让别人在任何厨房都能做出一模一样的味道。

训练模型时用的框架（比如PyTorch、TensorFlow这些）就像你熟悉的自家厨房，锅碗瓢盆怎么摆、火候怎么控你都门儿清，但真要开连锁店，就得把操作流程标准化，让别的厨师用普通厨具也能复现，模型导出就是这个“标准化”的过程。

为什么不能直接用训练时的文件？

好问题！训练过程中保存的检查点（checkpoint）文件确实包含了模型的全部信息——网络结构、权重参数、优化器状态等等，但这些东西太“原始”了，就像给你一堆食材和一本厚厚的烹饪笔记,要做菜还得从头研究。

训练框架通常比较“笨重”，依赖一大堆库和环境配置，你想啊，如果每部署一个模型都得安装整个训练环境，那得多麻烦？导出就是为了让模型能轻装上阵,独立运行。

常见的模型格式有哪些？

ONNX：这玩意儿像个万能翻译器，不管你在哪个框架下训练的模型，转换成ONNX格式后，就能在各种硬件和平台上运行,特别适合需要在不同环境中迁移的场合。
TensorFlow SavedModel：如果你用TF训练，这是它的“标准行李打包法”，会把模型的计算图、权重、还有必要的资源文件一起打包,非常完整。
TorchScript：PyTorch用户的打包方案，可以把动态的PyTorch模型转换成静态图，提高运行效率,特别适合生产环境。
TFLite：专门为移动设备和嵌入式设备设计的轻量级格式，就像把满汉全席做成方便食品,既保留风味又方便携带。
Core ML：苹果全家桶用户的专属格式，想在iPhone、iPad上跑AI？用这个就对了。

实操一下：导出流程长啥样？

以PyTorch为例,导出ONNX格式其实就几行代码：

import torch
import torchvision
model = torchvision.models.resnet18(pretrained=True)
model.eval() # 切换到推理模式
# 随便搞个模拟输入
dummy_input = torch.randn(1, 3, 224, 224)
# 开导！
torch.onnx.export(model, dummy_input, "my_resnet.onnx", 
                 input_names=['input'], output_names=['output'],
                 dynamic_axes={'input': {0: 'batch_size'}, 
                              'output': {0: 'batch_size'}})

看，关键就是那个torch.onnx.export函数，需要注意的是，你得给一个示例输入（就是那个dummy_input），让系统知道你的数据长什么样。dynamic_axes参数是用来处理可变批次大小的——毕竟实际使用时,一次处理1张图片和一次处理100张图片都很常见。

踩坑预警：这些雷区别碰！

新手导出模型时,经常会遇到这些幺蛾子：

张量形状不匹配：训练时输入可能是[32, 3, 224, 224]（32张224x224的RGB图片），但推理时变成了[1, 3, 256, 256]，如果模型里有写死的形状定义,那就悲剧了。
动态控制流丢失：如果你的模型里有“如果A大于B就执行X否则执行Y”这样的逻辑，转换成静态图时可能会出问题,ONNX对这些动态操作的支持还在不断完善。
自定义操作不支持：你自个儿写的特殊层或者操作，可能不在ONNX的标准操作集里，这时候就需要自己实现自定义操作符,或者想办法用现有操作组合出来。
后处理没打包：模型输出的可能是一堆原始数据，还需要经过非极大值抑制（NMS）之类的后处理才能变成最终结果,别忘了把这些后处理步骤也打包进去。

导出后别忘了验证！

模型导出了不等于万事大吉,一定要验证导出的模型和原始模型行为是否一致：

import onnxruntime
import numpy as np
# 用原始模型推理
original_output = model(dummy_input).detach().numpy()
# 用导出的ONNX模型推理
ort_session = onnxruntime.InferenceSession("my_resnet.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_output = ort_session.run(None, ort_inputs)
# 对比结果
np.testing.assert_allclose(original_output, ort_output[0], rtol=1e-03, atol=1e-05)
print("恭喜，模型导出成功！")

如果两个输出差不多（在允许的误差范围内）,那就可以开香槟庆祝了！

打包好了，然后呢？

导出的模型就像打包好的行李,接下来可以去任何地方：