如何保存TensorFlow模型并在其他文件或Colab重启后加载使用？

先排查可能的操作错误 #

• 路径问题：Colab的临时目录在会话断开后会被清空，直接保存到默认路径的模型会丢失，必须保存到挂载的Google Drive；跨文件加载时，确保加载代码的工作目录和模型保存路径一致，或者使用绝对路径。
• 版本兼容：保存模型的TensorFlow版本和加载时的版本差异过大（比如跨大版本）会导致加载失败，尽量保持版本一致，至少保证大版本相同（比如都是TF2.x）。
• 自定义层/结构问题：如果模型包含自定义层，加载前必须先定义好对应的自定义层类，否则TensorFlow无法解析模型结构；另外要确保保存前模型已经完成训练、结构完全确定（比如没有未编译的层）。

可行的替代保存加载方法 #

1. SavedModel格式（官方推荐） #

这是TensorFlow的标准序列化格式，兼容性最好，支持跨环境和部署场景：

# 保存模型
model.save("my_saved_model")  # 会生成一个文件夹，包含模型结构和权重

# 加载模型
import tensorflow as tf
loaded_model = tf.keras.models.load_model("my_saved_model")

2. 分离保存结构与权重 #

适合需要单独调整模型结构或复用权重的场景，也能避免部分序列化问题：

# 保存模型结构（JSON格式）
model_json = model.to_json()
with open("model_structure.json", "w") as json_file:
    json_file.write(model_json)
# 保存权重
model.save_weights("model_weights.h5")

# 加载模型
from tensorflow.keras.models import model_from_json
# 先加载结构
with open("model_structure.json", "r") as json_file:
    loaded_model_json = json_file.read()
loaded_model = model_from_json(loaded_model_json)
# 再加载权重
loaded_model.load_weights("model_weights.h5")
# 注意：加载后需要重新编译模型（如果要继续训练或评估）
loaded_model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

3. Colab环境专用方案 #

在Colab中必须挂载Google Drive来持久化模型，避免会话断开丢失：

# 先挂载Google Drive
from google.colab import drive
drive.mount('/content/drive')

# 保存模型到Drive
model.save("/content/drive/MyDrive/trained_models/my_model.keras")

# 加载模型（会话重启后）
import tensorflow as tf
loaded_model = tf.keras.models.load_model("/content/drive/MyDrive/trained_models/my_model.keras")

内容的提问来源于stack exchange，提问作者rand1