防止在训练模型时信息丢失 用于TensorFlowKeras和PyTorch的检查点教程

Posted 2021-05-02 ATYUN订阅号

AiTechYun

编辑：Yining

如果你玩过电子游戏，你就会明白为什么检查点（chekpoint）是有用的了。举个例子，有时候你会在一个大Boss的城堡前把你的游戏的当前进度保存起来——以防进入城堡里面就Game Over了。

机器学习和深度学习实验中的检查点本质上是一样的，它们都是一种保存你实验状态的方法，这样你就可以从你离开的地方开始继续学习。

如果你因为停电、操作系统故障、工作优先或其他类型的意外错误而丢失了一个或多个实验，你一定会抓狂。其他时候，即使你没有遇到不可预见的错误，你也可能只是想要恢复一种新实验的训练的特殊状态，或者从一个给定的状态中尝试不同的事情。

这就是为什么你需要检查点！

但是，等等，还有一个很重要的原因。如果你在工作结束时不检查你的训练模式，你将会失去所有的结果！简单来说，如果你想使用你训练的模型，你就需要一些检查点。

FloydHub是一个极其易用的深度学习云计算平台。号称“Zero Setup for Deep Learning”。它的服务主旨是： “您就专心研究您的深度学习，其它的环境配置、部署、版本控制等等都交给我们来做就可以了”。

• FloydHub网址：https://www.floydhub.com

这篇文章将演示如何在FloydHub上对你的训练模型进行检查，以便你可以从这些保存的状态重新开始你的实验。

什么是检查点?

Keras文档为检查点提供了一个很好的解释:

• 模型的体系结构，允许你重新创建模型

• 模型的权重

• 训练配置(损失、优化器、epochs和其他元信息)

• 优化器的状态，允许在你离开的地方恢复训练

同样，一个检查点包含了保存当前实验状态所需的信息，以便你可以从这一点恢复训练。

检查点策略

你可以根据你正在执行的训练类型，采用不同的检查点策略。

• 短期训练制度(几分钟到几小时)

• 正常的训练制度(数小时到一整天)

• 长期训练制度(数天至数周)

短期训练制度
典型的做法是在训练结束时，或者在每个epoch结束时，保存一个检查点。

正常的训练制度
在这种情况下，在每个n_epochs中保存多个检查点，并跟踪我们所关心的一些验证度量，这是很常见的。通常，有一个固定的最大数量的检查点，这样就不会占用太多的磁盘空间(例如，将你最大的检查点数量限制在10个，新的位置将会取代最早的检查点)。

长期训练制度
在这种类型的训练体系中，你可能希望采用与常规机制类似的策略：在每一个n_epochs中，你都可以节省多个检查点，并在你所关心的验证度量上保持最佳状态。在这种情况下，由于训练将花费很长的时间，所以减少检查点的次数是很常见的，但是需要维护更多的检查点。

哪种制度适合我?

这些不同策略之间的折衷是要保持频率和检查点文件的数量。让我们来看看当我们对这两个参数进行操作时发生了什么：

在FloydHub中保存和恢复

现在，让我们研究FloydHub上的一些代码。我将向你展示如何在TensorFlow、Keras和PyTorch这三个流行的深度学习框架中保存检查点:

在开始之前，使用floyd login命令登录到FloydHub命令行工具，然后复刻（fork）并初始化（init）项目:

$ git clone https://github.com/floydhub/save-and-resume.git
$ cd save-and-resume
$ floyd init save-and-resume

对于我们的检查点示例，我们将使用深度学习的Hello，World:使用卷积神经网络模型的MNIST分类任务。

因为预先清楚我们的检查点策略是很重要的，我将说明我们将要采用的方法:

• 只保留一个检查点

• 在每个epoch结束时采取策略

• 保存具有最佳(最大)验证精确度的那个

如果是这样的小例子，我们可以采用短期的训练制度。

命令
在深入研究具体的工作示例之前，让我们概述一下你需要的基本命令。当你开始新工作时，你的第一个命令看起来是这样的:

floyd run \
 [--gpu] \
 --env <env> \
 --data <your_dataset>:<mounting_point_dataset> \
 "python <script_and_parameters>"

提示:在你的python脚本中，你需要确保将检查点保存到/output文件夹中。FloydHub将自动保存/outputdirectory的内容作为工作的输出，这就是你将如何利用这些检查点来恢复工作的方式。

一旦你的工作完成，你就可以将该工作的输出作为下一项工作的输入进行挂载（mount），从而允许你的脚本利用你在该项目的下一个运行中创建的检查点。

floyd run \
 [--gpu] \
 --env <env> \
 --data <your_dataset>:<mounting_point_dataset> \
 --data <output_of_previous_job>:<mounting_point_model> \
 "python <script_and_parameters>"

好了,让我们看看如何使用这三个框架来实现这一点。

TensorFlow

• 在FloydHub Jupyter Notebook查看完整例子：https://www.floydhub.com/redeipirati/projects/save-and-resume/53/code/tf_mnist_cnn_jupyter.ipynb

TensorFlow提供了不同的保存和恢复检查点的方法。在我们的例子中，我们将使用tf.Estimator API，这种API背后采用了tf.train.Saver,tf.train.CheckpointSaverHook和tf.saved_model.builder.SavedModelBuilder三种函数。

更详细地说，tf.EstimatorAPI使用第一个函数来保存检查点，第二个函数根据所采用的检查点策略进行操作，最后一个以使用export_savedmodel()方法导出模型。

保存一个TensorFlow检查点
在初始化一个评估器之前，我们必须定义检查点策略。为此，我们必须使用tf.estimator.RunConfig API为预估程序创建一个配置。这里有一个例子，我们可以这样做:

# Save the checkpoint in the /output folder
filepath = "/output/mnist_convnet_model"

# Checkpoint Strategy configuration
run_config = tf.contrib.learn.RunConfig(
 model_dir=filepath,
 keep_checkpoint_max=1)

通过这种方式，我们告诉预估者应该从哪个目录中保存或恢复一个检查点，以及要保存多少个检查点。

接下来，我们必须在评估器的初始化中提供这个配置:

# Create the Estimator
mnist_classifier = tf.estimator.Estimator(
 model_fn=cnn_model_fn, config=run_config)

现在我们已经设置好了在TensorFlow代码中保存检查点。

恢复一个TensorFlow检查点
我们也已经准备好从下一个实验运行的检查点重新开始。如果评估器在给定的模型文件夹中找到一个检查点，那么它将从最后一个检查点加载。

下面是运行TensorFlow检查点示例的步骤。

通过FloydHub命令模式
第一次训练命令:

floyd run \
 --gpu \
 --env tensorflow-1.3 \
 --data redeipirati/datasets/mnist/1:input \
 'python tf_mnist_cnn.py'

• –env标记指定该项目应该运行的环境(在Python3.0.6上的Tensorflow 1.3.0 + Keras 2.0.6)

• --data标记指定pytorch-mnist数据集应该在/inputdirectory中可以使用

• –gpu标记实际上是可选的，除非你想马上开始运行GPU机器上的代码

从你的检查点恢复：

floyd run \
 --gpu \
 --env tensorflow-1.3 \
 --data redeipirati/datasets/mnist/1:input \
 --data <your-username>/projects/save-and-resume/<jobs>/output:/model \
 'python tf_mnist_cnn.py'

• –env标记指定该项目应该运行的环境(在Python3.0.6上的Tensorflow 1.3.0 + Keras 2.0.6)

• 第一个 --data标记指定pytorch-mnist数据集应该在/inputdirectory中可以使用

• 第二个–data标记指定前一个工作的输出应该在/modeldirectory中可以使用

• –gpu标记实际上是可选的——除非你想马上开始运行GPU机器上的代码

通过FloydHub的Jupyter Notebook模式

floyd run \
 --gpu \
 --env tensorflow-1.3 \
 --data redeipirati/datasets/mnist/1:input \
 --mode jupyter

• –env标记指定该项目应该运行的环境(在Python3.0.6上的Tensorflow 1.3.0 + Keras 2.0.6)

• –data标记指定pytorch-mnist数据集应该在/inputdirectory中可以使用

• –gpu标记实际上是可选的——除非你想马上开始运行GPU机器上的代码

• –mode标记指定该工作应该提供一个Jupyter notebook实例

从你的检查点恢复:
如果你想从你的Jupyter notebook上的前一份工作中加载一个检查点，那么只需添加–data <your-username>/projects/save-and-resume/<jobs>/output:/model到之前的命令。

Keras

• 在FloydHub Jupyter Notebook查看完整例子：https://www.floydhub.com/redeipirati/projects/save-and-resume/53/code/keras_mnist_cnn_jupyter.ipynb

Keras为保存和加载检查点提供了一个很棒的API。让我们来看看:

保存一个Keras检查点
Keras提供了一组名为回调（callbacks）的函数:你可以把回调看作是在某些训练状态下触发的事件。我们需要用于检查点的回调是ModelCheckpoint，它根据我们在示例中采用的检查点策略提供所需的所有特性。

注意:这个函数只会保存模型的权重——如果你想保存整个模型或部分组件，你可以在保存模型时查看Keras文档。

首先，我们必须导入回调函数:

from keras.callbacks import ModelCheckpoint

接下来，就在对model.fit(…)的调用之前，是时候准备检查点策略了。

# Save the checkpoint in the /output folder
filepath = "/output/mnist-cnn-best.hdf5"

# Keep only a single checkpoint, the best over test accuracy.
checkpoint = ModelCheckpoint(filepath,
 monitor='val_acc',
 verbose=1,
 save_best_only=True,
 mode='max')

• filepath=”/output/mnist-cnn-best.hdf5″:记住，FloydHub将保存/output文件夹的内容！

• monitor=’val_acc’:这是我们所关心的度量：验证精确度

• verbose=1:它将打印更多信息

• save_best_only=True:只保留最好的检查点(在最大化验证精确度的情况下)

• mode=’max’:以最大化验证精确度保存检查点

默认情况下，周期(或检查点频率)设置为1，这意味着在每个epoch结束。

最后，我们已经准备好看到在模型训练期间应用的检查点策略。为了做到这一点，我们需要将回调变量传递给model.fit(…)调用:

# Train
model.fit(x_train, y_train,
 batch_size=batch_size,
 epochs=epochs,
 verbose=1,
 validation_data=(x_test, y_test),
 callbacks=[checkpoint]) # <- Apply our checkpoint strategy

根据我们选择的策略，你会看到:

# This line when the training reach a new max
Epoch < n_epoch >: val_acc improved from < previous val_acc > to < new max val_acc >, saving model to /output/mnist-cnn-best.hdf5

# Or this line
Epoch < n_epoch >: val_acc did not improve

就是这样，你现在已经设置好保存你的检查点了。

恢复一个Keras检查点
Keras模型提供了load_weights()方法，该方法从hdf5file文件中加载权重。要加载模型的权重，你只需在模型定义之后添加这一命令行:

... # Model Definition

model.load_weights(resume_weights)

下面是如何在FloydHub运行这个Keras的示例:

通过FloydHub的命令模式
第一次训练命令:

floyd run \
 --gpu \
 --env tensorflow-1.3 \
 'python keras_mnist_cnn.py'

• –env标记指定该项目应该运行的环境(在Python3.0.6上的Tensorflow 1.3.0 + Keras 2.0.6)

• –gpu标记实际上是可选的——除非你想马上开始运行GPU机器上的代码

Keras提供了一个用于处理MNIST数据的API，因此我们可以在本例中跳过数据集的安装。

从你的检查点恢复：

floyd run \
 --gpu \
 --env tensorflow-1.3 \
 --data <your-username>/projects/save-and-resume/<jobs>/output:/model \
 'python keras_mnist_cnn.py'

• –env标记指定该项目应该运行的环境(在Python3.0.6上的Tensorflow 1.3.0 + Keras 2.0.6)

• --data标记指定之前工作的输出应该在/modeldirectory中可以使用

• –gpu标记实际上是可选的——除非你想马上开始运行GPU机器上的代码

通过FloydHub的Jupyter Notebook模式

floyd run \
 --gpu \
 --env tensorflow-1.3 \
 --mode jupyter

• –env标记指定该项目应该运行的环境(在Python3.0.6上的Tensorflow 1.3.0 + Keras 2.0.6)

• –gpu标记实际上是可选的——除非你想马上开始运行GPU机器上的代码

• –mode标记指定该工作应该提供一个Jupyter notebook实例

从你的检查点恢复:
如果你想要从以前的工作中加载一个检查点，那么只需添加–data <your-username>/projects/save-and-resume/<jobs>/output:/model。

PyTorch

• 在FloydHub Jupyter Notebook查看完整例子：https://www.floydhub.com/redeipirati/projects/save-and-resume/53/code/pytorch_mnist_cnn_jupyter.ipynb

不幸的是，目前，检查点对于PyTorch的API来说并不像Keras那样容易。我们需要根据所选的检查点策略编写自己的解决方案。

保存一个PyTorch检查点
PyTorch没有提供一个一体化（all-in-one）的API来定义一个检查点策略，但是它提供了一个简单的方法来保存和恢复一个检查点。根据语义序列化（semantic serialization）的官方文档，最好的做法是只保存权重，这是由于代码重构问题造成的。

• 语义序列化文档：http://pytorch.org/docs/master/notes/serialization.html

因此，让我们来看看如何在PyTorch中保存模型的权重。

首先，让我们定义一个save_checkpoint函数，该函数负责处理要保留的检查点数量和文件序列化的所有指令。

def save_checkpoint(state, is_best, filename='/output/checkpoint.pth.tar'):
 """Save checkpoint if a new best is achieved"""
 if is_best:
 print ("=> Saving a new best")
 torch.save(state, filename) # save checkpoint
 else:
 print ("=> Validation Accuracy did not improve")

然后，在训练中(通常是一个循环的次数)，我们定义了检查点的频率(在我们的例子中，指的是在每个epoch结束时)和我们想要存储的信息(epoch，模型的权重，以及达到的最佳精确度):

...

# Training the Model
for epoch in range(num_epochs):
 train(...) # Train
 acc = eval(...) # Evaluate after every epoch

# Some stuff with acc(accuracy)
 ...

# Get bool not ByteTensor
 is_best = bool(acc.numpy() > best_accuracy.numpy())
 # Get greater Tensor to keep track best acc
 best_accuracy = torch.FloatTensor(max(acc.numpy(), best_accuracy.numpy()))
 # Save checkpoint if is a new best
 save_checkpoint({
 'epoch': start_epoch + epoch + 1,
 'state_dict': model.state_dict(),
 'best_accuracy': best_accuracy
 }, is_best)

你现在可以在你的PyTorch实验中保存检查点。

恢复一个PyTorch检查点
为了恢复一个PyTorch检查点，我们必须在训练前加载我们需要的权重和元信息。

# cuda = torch.cuda.is_available()
if cuda:
 checkpoint = torch.load(resume_weights)
else:
 # Load GPU model on CPU
 checkpoint = torch.load(resume_weights,
 map_location=lambda storage,
 loc: storage)
start_epoch = checkpoint['epoch']
best_accuracy = checkpoint['best_accuracy']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (trained for {} epochs)".format(resume_weights, checkpoint['epoch']))

下面是如何在FloydHub运行这个PyTorch的示例:

通过FloydHub的命令模式
第一次训练命令:

floyd run \
 --gpu \
 --env pytorch-0.2 \
 --data redeipirati/datasets/pytorch-mnist/1:input \
 'python pytorch_mnist_cnn.py'

• –env标记指定该项目应该运行的环境(在Python 3上的PyTorch 0.2.0)

• –data标记指定pytorch-mnist数据集应该在/inputdirectory中可用

• –gpu标记实际上是可选的——除非你想马上开始运行GPU机器上的代码

从你的检查点恢复：

floyd run \
 --gpu \
 --env pytorch-0.2 \
 --data redeipirati/datasets/pytorch-mnist/1:input \
 --data <your-username>/projects/save-and-resume/<jobs>/output:/model \
 'python pytorch_mnist_cnn.py'

• –env标记指定该项目应该运行的环境(在Python 3上的PyTorch 0.2.0)

• 第一个–data标记指定pytorch-mnist数据集应该在/inputdirectory中可以使用

• 第二个–data标记指定前一个工作的输出应该在/modeldirectory中可以使用

• –gpu标记实际上是可选的——除非你想马上开始运行GPU机器上的代码

通过FloydHub的Jupyter Notebook模式

floyd run \
 --gpu \
 --env pytorch-0.2 \
 --data redeipirati/datasets/pytorch-mnist/1:input \
 --mode jupyter

• –env标记指定该项目应该运行的环境(在Python 3上的PyTorch 0.2.0)

• –data标记指定pytorch-mnist数据集应该在/inputdirectory中可以使用

• –gpu标记实际上是可选的——除非你想马上开始运行GPU机器上的代码

• –mode标记指定该工作应该为我们提供一个Jupyter notebook

从你的检查点恢复:
如果你想要从以前的工作中加载一个检查点，那么只需添加–data <your-username>/projects/save-and-resume/<jobs>/output:/model。