手把手教你导入训练好的AI模型，别让心血卡在最后一步！

搞AI模型训练的朋友，估计都经历过这种时刻：辛辛苦苦调参、跑数据，模型好不容易训出来了，性能指标看着也挺美，结果一到“导入”这步，懵了，一堆文件，各种格式，该拖哪儿？点哪个按钮？命令行报错看得人头大,感觉几天几夜的心血就要卡在这临门一脚上了。

别急，这种事儿太常见了，今天咱就抛开那些复杂晦涩的理论，像朋友聊天一样，聊聊怎么把你这“宝贝模型”顺顺当当地“请”到该去的地方，让它真正开始干活儿，这里没有万能公式，因为不同的框架、不同的使用场景，方法天差地别,但核心思路是相通的。

你得知道你的“宝贝”打包成了什么样。

模型训练完，它可不是一个孤零零的文件就完事了（虽然有时确实是），你通常会得到一堆“遗产”：

• 检查点文件：TensorFlow 的 .ckpt 文件（可能是三个：.data, .index, .meta），这就像是游戏的存档点，包含了训练到某一刻的所有参数,方便你后续接着训或者恢复。
• 整个模型结构+权重：像 Keras 喜欢用的 .h5 文件，或者 PyTorch 常见的 .pth、.pt 文件，这通常是把模型的结构和训练好的权重打包在一起了,用起来比较省心。
• “冻结”后的模型：TensorFlow 的 .pb 文件，这相当于把模型结构和参数“冰封”成一个独立的、高效的计算图，特别适合部署到生产环境,比如移动端或服务器上。
• ONNX 格式：这是个“中间人”格式，目标是为了让不同框架训练的模型能互相交流，如果你想把 PyTorch 训练的模型拿到其他地方用，转成 .onnx 文件是个好主意。

导入的第一步不是急着操作，而是先看看你的输出目录里，到底生成了些啥，知道手里拿的是什么“钥匙”，才能去找对应的“锁”。

就是“对号入座”的环节了。

1. 如果你想接着训练或者做精细调整： 那你需要的是加载检查点，以 TensorFlow 2.x 为例,大概会是这样：

   # 假设你的模型结构已经定义好了，叫 my_model
   my_model = create_my_model_architecture() # 这是你定义模型结构的函数
   checkpoint_path = "./training_checkpoints/your_model.ckpt"
   # 只加载权重，不是整个模型
   my_model.load_weights(checkpoint_path)

   这样，之前训练好的参数就覆盖到你的模型结构上了，你可以继续喂数据训练，或者用它来做预测，PyTorch 的思路也类似，用 torch.load 加载 .pth 文件，然后是一个 model.load_state_dict() 的操作,把参数字典塞回模型里。

2. 如果你想快速用它做预测（推理）： 这时，加载整个模型文件通常最方便，比如用 Keras 保存的 .h5 文件：

   from tensorflow import keras
   loaded_model = keras.models.load_model("my_complete_model.h5")
   # 然后就可以直接用了
   predictions = loaded_model.predict(new_data)

   一步到位，结构和权重都有了，PyTorch 里直接 torch.load 整个模型对象也行，但要注意代码里模型类的定义必须和保存时一致，不然会找不到“家”。

3. 如果你要部署到生产环境： 这就轮到“冻结”模型或者转换格式出场了，TensorFlow 的 .pb 文件，你需要用专门的 API 来加载：

   import tensorflow as tf
   # 加载计算图
   with tf.io.gfile.GFile("frozen_model.pb", 'rb') as f:
       graph_def = tf.compat.v1.GraphDef()
       graph_def.ParseFromString(f.read())
   # 导入到当前默认的图中
   tf.import_graph_def(graph_def, name='')
   # 然后通过张量名称来获取输入输出，进行推理

   这个过程稍微绕一点，需要你清楚模型输入输出的节点名称，而 ONNX 格式的模型，则有对应的 onnxruntime 库来加载和运行，能脱离原始训练框架,效率很高。

几个容易栽跟头的小坑：

• 路径问题：老是说“文件找不到”？检查一下你的文件路径是相对路径还是绝对路径，字符串里有没有写对，在 Python 里，路径中的反斜杠 \ 最好改成正斜杠 或者在字符串前加 r 防止转义。
• 版本地狱：这是最头疼的！用 TensorFlow 1.x 保存的 .ckpt，拿到 TensorFlow 2.x 环境加载，很可能报错，PyTorch 的模型在不同版本间有时也会闹脾气，尽量保证保存和加载的环境一致，如果不行，就要查官方文档的版本兼容说明，或者考虑用 ONNX 这种中间格式过渡一下。
• 依赖缺失：模型里用了个自定义的层或者特殊的处理函数？保存的时候可能没打包进去，加载的时候，Python 环境里必须能找到完全一样的类或函数定义，否则就会提示你“不认识这个东西”,这时候你就得把当初定义它的代码也搬过来。

说到底，导入模型不是什么高深魔法，它就是个“按图索骥”的体力活加细心活，核心就是搞清楚格式，用对方法，注意环境，别被那些命令行错误吓到，大部分错误信息网上都能搜到解决方案,关键是把错误提示看懂。

下次再遇到模型导入的问题，深呼吸，先回来看看这篇文章，一步步核对，你的AI心血，值得被稳稳当当地接生到现实世界，去发挥它的价值,祝你一次导入成功！

（免费申请加入）AI工具导航网

相关标签： # 训练好的ai模型怎么导入