侧边栏壁纸
  • 累计撰写 244 篇文章
  • 累计创建 16 个标签
  • 累计收到 0 条评论
隐藏侧边栏

TensorFlow 2 模型保存恢复

kaixindeken
2021-09-24 / 0 评论 / 0 点赞 / 64 阅读 / 5,084 字

模型保存是深度学习中的重要一环,最常见的是两种情况。首先,在模型的训练过程中,可能会发生另其意外终止的各种情况,所以实时保存模型就非常重要了。我们可以在训练终止时快速恢复模型训练过程,不至于前功尽弃。

此外,深度学习可能需要利用训练好的模型进行推理预测,此时需要调用已经保存好的模型。面对这两种情形,一个完善的深度学习框架都会提供相应的 API 以保证能对模型进行存储。接下来,我们以 TensorFlow 常用的模型保存方法为例,来学习如果讲一个正在训练或已训练完成的模型和权重保存起来。

首先,我们以 TensorFlow Keras 训练模型为例,这里直接沿用前面实验中用过的 DIGITS 数据集及 Keras 顺序模型结构。

import numpy as np
import tensorflow as tf
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split

digits = load_digits()  # 加载数据集
digits_y = np.eye(10)[digits.target]  # 标签独热编码
# 切分数据集,测试集占 20%
X_train, X_test, y_train, y_test = train_test_split(digits.data, digits_y,
                                                    test_size=0.2, random_state=1)
# 输出训练特征,训练目标值,测试特征,测试目标值形状
X_train.shape, y_train.shape, X_test.shape, y_test.shape
model = tf.keras.Sequential()  # 建立顺序模型
model.add(tf.keras.layers.Dense(units=30, input_dim=64, activation='relu'))  # 隐含层
model.add(tf.keras.layers.Dense(units=10, activation='softmax'))  # 输出层
# 编译模型,添加优化器,损失函数和评估方法
model.compile(optimizer='adam', loss='categorical_crossentropy',
              metrics=['accuracy'])
model.summary()  # 查看模型结构

正常情况下,接下来就可以使用 model.fit 开始训练过程了。

保存模型权重

要想在训练过程中实时保存模型,需要借助于 TensorFlow 提供的 Checkpoint 检查点机制。调用 tf.keras.callbacks.ModelCheckpoint,可以在训练的过程中和结束时回调保存的模型。这样,我们就可以在任意时刻从打断的地方继续开始训练,以防止前功尽弃。

接下来,我们创建一个保存模型权重的回调:

import os

checkpoint_path = "save_model/digits.ckpt"  # 保存模型的路径和名称
checkpoint_dir = os.path.dirname(checkpoint_path)  # 新建一个单独的文件夹存放

# 创建一个保存模型权重的回调
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True)

沿用之前一致的训练代码,唯一的区别在于添加 callbacks 参数,使用新的回调训练模型。

# 使用新的回调训练模型
model.fit(X_train, y_train, batch_size=64, epochs=10,
          validation_data=(X_test, y_test), callbacks=[cp_callback])

训练的过程中,每一次迭代之后都会有 saving model to 的字眼,意味着模型的权重已被实时保存。此时,我们使用 Linux 命令来查看相应目录,你应该可以看到 3 个序列文件。

ls {checkpoint_dir}

要想恢复保存的模型权重非常简单,借助于 .load_weights 方法即可。

# 加载已保存的模型权重
model.load_weights(checkpoint_path)

此时,你就可以重新使用 model.fit 方法来继续完成训练,或者使用 model.evaluate 对模型进行评估了。

tf.keras.callbacks.ModelCheckpoint 调用时提供了其他的参数,例如可以设置指定个 Epoch 之后保存一次,这对于较大的模型训练过程很实用,毕竟保存一次需要花费很长的时间。

另外,我们可以在保存模型权重路径中加入 Epoch 标记,这样可以把每个 Epoch 的权重都保存下来,而不是只保存最新的一个。

# 添加 Epoch 标记信息 (使用 `str.format`)
checkpoint_path = "save_model/digits-{epoch:04d}.ckpt"  # 保存模型的路径和名称
# 创建一个保存模型权重的回调
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True)
# 重新使用新的回调训练模型
model.fit(X_train, y_train, batch_size=64, epochs=10,
          validation_data=(X_test, y_test), callbacks=[cp_callback])

此时,查看目录下方应该会有 10 个 Epoch 训练保存的模型。

ls {checkpoint_dir}

那么,在恢复权重时,你需要指定相应的权重文件名称。权重文件是以 .ckpt 结尾的文件。例如,恢复第 5 个 Epoch 训练时的权重。

# 恢复第 5 个 Epoch 训练时的权重
model.load_weights('save_model/digits-0005.ckpt')

保存完整模型

上面提到的 Checkpoint 检查点机制只能保存模型的权重,这些权重只能被相同结构的模型进行加载。也就是说,你无法将保存下来的权重恢复到其他不同的模型结构中并继续训练。

上面之所以能使用 model.load_weights,是因为 model 的结构已经定义好了,也就是和保存时一模一样的结构。那么,如果你想保存一个模型,而在下一次恢复该模型时无需考虑结构,则需要保存整个模型。

TensorFlow Keras 可以使用 HDF5 格式来保存整个模型,区别于仅保存模型权重,保存整个模型时包括:模型权重,模型结构,优化器配置信息。保存整个模型的方法很简单,只需要使用 model.save 即可。

# 将上方训练好的整个模型保存为 HDF5 文件
model.save('digits.h5')

接下来,我们使用 TensorFlow 来加载保存的完整模型,将其定义为新的模型。

# 加载保存的完整模型
new_model = tf.keras.models.load_model('digits.h5')

# 显示网络结构
new_model.summary()

此时可以看到已加载的模型结构,同时可以直接使用该模型。例如,对测试数据进行预测。

new_model.predict(X_test)

至此,你已经掌握了 TensorFlow Keras 的模型权重保存和完整模型保存方法。那么,你可能在想使用低阶 API 构建的模型该如何保存,当然会比 Keras 复杂一些。

低阶 API 保存权重

我们这里说的低阶 API 代之使用 tf.GradientTape() 去追踪梯度并优化的过程。不过模型建立时,依旧需要参考 Keras 混合模型的搭建方法,我们构建一模一样的模型。

# 继承 tf.keras.Model 构建模型类
class Model(tf.keras.Model):
    def __init__(self):
        super(Model, self).__init__()
        self.dense_1 = tf.keras.layers.Dense(30, activation='relu')  # 初始化
        self.dense_2 = tf.keras.layers.Dense(10, activation='softmax')

    def call(self, inputs):
        x = self.dense_1(inputs)  # 前向传播过程
        return self.dense_2(x)

接下来,参考前面简单神经网络实验中,实现损失函数和准确度计算函数。

def loss_fn(model, x, y):
    preds = model(x)
    return tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=preds, labels=y))
def accuracy_fn(logits, labels):
    preds = tf.argmax(logits, axis=1)  # 取值最大的索引,正好对应字符标签
    labels = tf.argmax(labels, axis=1)
    return tf.reduce_mean(tf.cast(tf.equal(preds, labels), tf.float32))

然后可以开始迭代过程。这里,首先使用 tf.train.Checkpoint 定义一个 Checkpoint 对象,传入使用的优化器和模型。然后,使用 tf.train.CheckpointManager 构建一个 Checkpoint 管理器对象,包含模型权重的保存目录,以及支持 max_to_keep 参数以保留最近的 N 个已保存权重。

EPOCHS = 100  # 迭代此时
LEARNING_RATE = 0.02  # 学习率

model = Model()  # 实例化模型类
opt = tf.optimizers.Adam(learning_rate=LEARNING_RATE)  # Adam 优化器

# 创建检查点对象
ckpt = tf.train.Checkpoint(step=tf.Variable(1), optimizer=opt, net=model)
# 创建检查点管理对象
manager = tf.train.CheckpointManager(ckpt, './save_model_1', max_to_keep=5)

# 使用判断,如果有保存的权重,则使用改权重,否则从头开始训练
ckpt.restore(manager.latest_checkpoint)
if manager.latest_checkpoint:
  print("找到已保存检查点,恢复训练 {}".format(manager.latest_checkpoint))
else:
  print("未找到已保存检查点,从头开始训练.")

for epoch in range(EPOCHS):
    with tf.GradientTape() as tape:  # 追踪梯度
        loss = loss_fn(model, X_train, y_train)

    trainable_variables = model.trainable_variables  # 需优化参数列表
    grads = tape.gradient(loss, trainable_variables)  # 计算梯度

    opt.apply_gradients(zip(grads, trainable_variables))  # 更新梯度
    accuracy = accuracy_fn(model(X_test), y_test)  # 计算准确度
    ckpt.step.assign_add(1)
    if int(ckpt.step) % 10 == 0:
        save_path = manager.save()
        print(f'Epoch [{epoch+1}/{EPOCHS}], Train loss: {loss}, Test accuracy: {accuracy}')
        print("Saved checkpoint for step {}: {}".format(int(ckpt.step), save_path))

上方,我们使用了一个判断语句决定是否会加载已保存的权重继续训练。当你第一次运行代码时,会从头开始训练。但当你重复执行上方的代码时,则会从已保存的最近权重处继续训练。

此外,当我们使用低阶 API 构建模型时,无法很方便地做到保存完整的模型。

0

评论区