3、TensorFlow模型基础：SavedModel格式、HDF5格式、Keras模型API、模型保存与加载

好，咱们进入正题。这一章聊的是TensorFlow模型的基础操作，说白了就是怎么把你的心血——训练好的模型——给存下来，以及怎么再把它请回来。你想想看，模型训练动辄几小时甚至几天，要是不会保存，那可就白干了。

我个人习惯把模型保存看作是「给AI拍身份证」。你得拍清楚，还得保证以后能认出来。TensorFlow里主要有两套身份证系统：一套是Keras自带的HDF5格式，另一套是官方的SavedModel格式。这两者有什么区别？什么时候该用哪个？嗯，这正是咱们要聊的。

3.1 Keras模型API：你每天都在用的老朋友

先说说Keras模型API。如果你用过TensorFlow，大概率已经接触过它了。Keras提供了三种构建模型的方式，我按自己的使用频率排个序：

  • Sequential顺序模型：最简单，一层层堆叠。适合大部分常规任务，比如图像分类、简单的回归。
  • Functional函数式API：灵活得多，支持多输入多输出、共享层。我在做多模态项目时几乎必用它。
  • Model子类化：最灵活，但也是最容易出坑的。你可以完全自定义前向传播逻辑。

举个例子，咱们用函数式API搭一个简单的模型：

import tensorflow as tf

# 定义输入
inputs = tf.keras.Input(shape=(784,))
# 定义隐藏层
x = tf.keras.layers.Dense(64, activation='relu')(inputs)
x = tf.keras.layers.Dense(64, activation='relu')(x)
# 定义输出
outputs = tf.keras.layers.Dense(10, activation='softmax')(x)

# 构建模型
model = tf.keras.Model(inputs=inputs, outputs=outputs)
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

你看，代码很直观。我个人建议初学者先从Sequential入手，等遇到复杂场景再切换到函数式API。至于子类化，嗯，除非你真的很清楚自己在做什么，否则慎用。我曾经在项目中用子类化写了一个自定义层，结果保存模型时折腾了一整天——后面会细说这个坑。

3.2 HDF5格式：Keras的经典保存方式

HDF5格式是Keras早期就支持的保存格式。它把模型的架构、权重、训练配置（比如优化器状态）全部打包到一个.h5文件里。保存和加载都非常简单：

# 保存模型
model.save('my_model.h5')

# 加载模型
loaded_model = tf.keras.models.load_model('my_model.h5')

就这么两行代码，搞定。但这里有个坑——HDF5格式不支持保存自定义层或自定义损失函数。如果你用了自定义的东西，加载时会报错。我记得有一次帮同事排查问题，他用了自定义的Focal Loss，保存成HDF5后怎么都加载不了。后来我让他改用SavedModel格式，问题就解决了。

⚠️ 注意：HDF5格式在加载自定义组件时，需要传入custom_objects参数。比如： 
model = tf.keras.models.load_model('my_model.h5', 
                                    custom_objects={'FocalLoss': FocalLoss})
如果你不想每次加载都传参，建议直接用SavedModel。

3.3 SavedModel格式：TensorFlow的官方标准

SavedModel是TensorFlow官方推荐的格式，也是模型部署时的首选。它不只是一个文件，而是一个目录，里面包含了：

  • saved_model.pb：模型的计算图定义
  • variables/：模型权重文件
  • assets/：额外的资源文件（比如词汇表）

保存和加载的代码也很简单：

# 保存为SavedModel格式
model.save('my_saved_model')

# 加载SavedModel
loaded_model = tf.keras.models.load_model('my_saved_model')

你可能会问：「那SavedModel比HDF5好在哪？」我总结一下：

特性 HDF5格式 SavedModel格式
文件结构 单个.h5文件 目录（含多个文件）
自定义组件支持 需额外处理 原生支持
部署兼容性 一般 优秀（TF Serving、TFLite等）
版本管理 不支持 支持（可保存多个版本）
跨语言调用 仅Python Python、C++、Java等

说白了，如果你只是做实验、写论文，HDF5完全够用。但如果你要部署到生产环境，或者要用TF Serving做推理服务，那SavedModel是唯一的选择。

3.4 模型保存与加载的实战技巧

光知道怎么保存还不够，咱们得聊聊实战中那些容易踩的坑。

3.4.1 只保存权重 vs 保存整个模型

有时候你不需要保存整个模型，只需要权重。比如你想用预训练模型做迁移学习，或者想在不同框架间转换。这时候可以用：

# 只保存权重
model.save_weights('model_weights.h5')

# 加载权重（需要先构建模型架构）
model.load_weights('model_weights.h5')

注意：加载权重前，你必须先定义好一模一样的模型架构。否则会报错。我曾经在项目中犯过这个错——改了模型结构忘了改加载代码，结果调试了半天才发现是权重形状不匹配。

3.4.2 训练过程中的自动保存

训练大模型时，我强烈建议使用回调函数自动保存。这样即使训练中途断了，也能从最近的检查点恢复：

checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
    filepath='checkpoint_{epoch:02d}.h5',
    save_weights_only=False,
    save_best_only=True,
    monitor='val_loss',
    mode='min'
)

model.fit(x_train, y_train, 
          epochs=50, 
          validation_data=(x_val, y_val),
          callbacks=[checkpoint_callback])

这里有个小技巧：save_best_only=True配合monitor='val_loss'，只会保存验证集损失最小的那个模型。这样你就不用手动筛选了。

3.4.3 自定义层的保存陷阱

嗯，这里要重点说一下。如果你用了自定义层，比如：

class MyLayer(tf.keras.layers.Layer):
    def __init__(self, units=32):
        super(MyLayer, self).__init__()
        self.units = units
    
    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b

保存模型时，TensorFlow需要知道这个自定义层怎么序列化。你需要在类里实现get_config()方法：

def get_config(self):
    return {'units': self.units}

如果不实现，加载模型时会报错：「Unknown layer: MyLayer」。我曾经在项目里被这个坑过，后来养成了习惯——所有自定义层都加上get_config()。

💡 小提示：如果你用的是子类化模型（Model子类化），保存时最好用SavedModel格式。HDF5对子类化模型的支持不太好，有时候能保存但加载不了。我个人的经验是：子类化模型一律用SavedModel，省心。

3.5 总结一下

这一章咱们聊了TensorFlow模型保存与加载的核心内容。说白了就是三件事：

  1. Keras模型API：Sequential、Functional、子类化，按需选择
  2. HDF5格式：简单方便，适合实验和迁移学习
  3. SavedModel格式：官方标准，适合部署和生产环境

下一章咱们会深入模型转换的核心——如何把TensorFlow模型转成ONNX格式。到时候你会看到，模型保存的格式选择直接影响到转换的成败。嗯，做好准备，咱们要开始动真格的了。

← 上一章 📖 返回目录 下一章 →