TensorFlow怎么保存训练好的权重_Python只保存模型参数而非全图

2026-06-02 Python

PyTorch中只保存权重应使用torch.save(model.state_dict(), path)，而非torch.save(model, path)；后者保存完整模型（结构+权重+优化器状态），体积大、有安全风险且版本兼容性差；加载权重前必须先实例化结构完全一致的模型，并调用load_state_dict()，推荐始终使用weights_only=True确保安全。

只保存权重用 `model.save_weights()`，不是 `model.save()`

很多人误以为 model.save() 是“保存模型”，其实它默认保存的是完整模型（架构 + 权重 + 优化器状态），文件体积大、加载时依赖原始代码定义。如果只想存参数（比如做模型蒸馏、迁移学习、或部署时分离权重和结构），必须显式调用 model.save_weights()。

常见错误现象：model.save('my_model.h5') 后试图用 tf.keras.models.load_model() 加载失败，报错 ValueError: No model found in config file 或无法重建自定义层——因为那根本不是权重文件，而是全模型 HDF5。

model.save_weights('weights.h5')：生成纯权重 HDF5，兼容性好，适合 TensorFlow 2.x 通用场景
model.save_weights('weights')（无后缀）：会自动保存为 TensorFlow 检查点格式（checkpoint + data-00000-of-00001），更轻量，但加载时需先有相同结构的模型实例
不推荐用 .ckpt 手动拼后缀——TensorFlow 会自动管理，硬写可能引发路径解析异常

加载权重前必须先构建相同结构的模型

model.load_weights() 不恢复架构，只填充变量值。这意味着你不能直接从权重文件“复活”一个模型——必须先用完全一致的代码（层类型、顺序、输入形状、自定义类注册）创建空模型，再调用 load_weights()。

典型踩坑场景：

训练时用了自定义 tf.keras.layers.Layer，加载时没 import 对应类 → 报错 Unknown layer
训练模型输入 shape 是 (32, 32, 3)，加载时新建模型设成 (224, 224, 3) → 权重张量 shape 不匹配，报 Shape mismatch
用 Sequential 训练，但加载时改用 Functional API 构建 → 层名/顺序不同，权重映射失败

验证方式：加载后执行 model.weights，检查数量和 shape 是否与训练模型一致。

`save_weights_only=True` 在 `ModelCheckpoint` 中才生效

训练中自动保存最佳权重，靠的是 tf.keras.callbacks.ModelCheckpoint。注意：它的 save_weights_only 参数只在回调里起作用，和 model.save_weights() 是两套机制。

Python 3.14.3

微软官方的 Python 扩展，是 VS Code 安装量最高的扩展（209M+）。集成 IntelliSense（通过 Pylance）、调试（通过 Python Debugger）、代码检查、格式化、重构和单元测试等功能。支持 Jupyter Notebook、虚拟环境管理和多 Python 版本切换。

正确写法：

cp_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath='best_weights.h5',
    save_weights_only=True,  # ← 关键：设为 True 才只存权重
    monitor='val_loss',
    save_best_only=True
)

容易忽略的点：

如果 filepath 写成 'best_model.h5' 却设 save_weights_only=True，结果仍是权重文件，但名字误导人
save_freq='epoch' 时，每个 epoch 都覆盖同名文件；想保留多版本，得用格式化字符串如 'weights_ep{epoch:02d}.h5'
检查点路径不含扩展名（如 'ckpt'）时，TensorFlow 默认用检查点格式；加了 .h5 才强制走 HDF5