模型加载及恢复训练

    在快速入门中,我们已经介绍了将训练好的模型保存到磁盘文件的方法。应用程序可以随时加载模型,完成预测任务。但是在日常训练工作中我们会遇到一些突发情况,导致训练过程主动或被动的中断。如果训练一个模型需要花费几天的训练时间,中断后从初始状态重新训练是不可接受的。

    万幸的是,飞桨支持从上一次保存状态开始训练,只要我们随时保存训练过程中的模型状态,就不用从初始状态重新训练。

    下面介绍恢复训练的代码实现,依然使用手写数字识别的案例,在网络定义的部分保持不变。

    1. loading mnist dataset from ./work/mnist.json.gz ......

    在开始介绍使用飞桨恢复训练前,先正常训练一个模型,优化器使用Adam,使用动态变化的学习率,学习率从0.01衰减到0.001。每训练一轮后保存一次模型,之后将采用其中某一轮的模型参数进行恢复训练,验证一次性训练和中断再恢复训练的模型表现是否一致(训练loss的变化)。

    注意进行恢复训练的程序不仅要保存模型参数,还要保存优化器参数。这是因为某些优化器含有一些随着训练过程变换的参数,例如Adam, Adagrad等优化器采用可变学习率的策略,随着训练进行会逐渐减少学习率。这些优化器的参数对于恢复训练至关重要。

    为了演示这个特性,下面训练程序使用adam优化器,学习率以多项式曲线从0.01衰减到0.001(polynomial decay)。

    • learning_rate:初始学习率
    • decay_steps:衰减步数
    • end_learning_rate:最终学习率
    • power:多项式的幂,默认值为1.0
    1. use_gpu = False
    2. place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace()
    3. with fluid.dygraph.guard(place):
    4. model = MNIST("mnist")
    5. model.train()
    6. EPOCH_NUM = 5
    7. BATCH_SIZE = 100
    8. # 定义学习率,并加载优化器参数到模型中
    9. total_steps = (int(60000//BATCH_SIZE) + 1) * EPOCH_NUM
    10. lr = fluid.dygraph.PolynomialDecay(0.01, total_steps, 0.001)
    11. # 使用Adam优化器
    12. for batch_id, data in enumerate(train_loader()):
    13. #准备数据,变得更加简洁
    14. image_data, label_data = data
    15. image = fluid.dygraph.to_variable(image_data)
    16. label = fluid.dygraph.to_variable(label_data)
    17. #前向计算的过程,同时拿到模型输出值和分类准确率
    18. predict, acc = model(image, label)
    19. avg_acc = fluid.layers.mean(acc)
    20. #计算损失,取一个批次样本损失的平均值
    21. loss = fluid.layers.cross_entropy(predict, label)
    22. avg_loss = fluid.layers.mean(loss)
    23. #每训练了200批次的数据,打印下当前Loss的情况
    24. if batch_id % 200 == 0:
    25. #后向传播,更新参数的过程
    26. avg_loss.backward()
    27. model.clear_gradients()
    28. # 保存模型参数和优化器的参数
    29. fluid.save_dygraph(model.state_dict(), './checkpoint/mnist_epoch{}'.format(epoch_id))
    30. fluid.save_dygraph(optimizer.state_dict(), './checkpoint/mnist_epoch{}'.format(epoch_id))

    恢复训练

    在上述训练代码中,我们训练了五轮(epoch)。在每轮结束时,我们均保存了模型参数和优化器相关的参数。

    • 使用model.state_dict()获取模型参数。
    • 使用optimizer.state_dict()获取优化器和学习率相关的参数。
    • 调用fluid.save_dygraph()将参数保存到本地。

    比如第一轮训练保存的文件是mnist_epoch0.pdparams,mnist_epoch0.pdopt,分别存储了模型参数和优化器参数。

    当加载模型时,如果模型参数文件和优化器参数文件是相同的,我们可以使用load_dygraph同时加载这两个文件,如下代码所示。

    1. params_dict, opt_dict = fluid.load_dygraph(params_path)

    如果模型参数文件和优化器参数文件的名字不同,需要调用两次load_dygraph分别获得模型参数和优化器参数。

    理想的恢复训练是模型状态回到训练中断的时刻,恢复训练之后的梯度更新走向是和恢复训练前的梯度走向是完全相同的。基于此,我们可以通过恢复训练后的损失变化,判断上述方法是否能准确的恢复训练。即从epoch 0结束时保存的模型参数和优化器状态恢复训练,校验其后训练的损失变化(epoch 1)是否和不中断时的训练完全一致。


    说明:

    恢复训练有如下两个要点:

    • 保存模型时同时保存模型参数和优化器参数。
    • 恢复参数时同时恢复模型参数和优化器参数。

    下面的代码将展示恢复训练的过程,并验证恢复训练是否成功。其中,我们重新定义一个train_again()训练函数,加载模型参数并从第一个epoch开始训练,以便读者可以校验恢复训练后的损失变化。

    1. epoch: 1, batch: 0, loss is: [0.04537241], acc is [0.98]
    2. epoch: 1, batch: 200, loss is: [0.04753707], acc is [0.98]
    3. epoch: 1, batch: 400, loss is: [0.02443215], acc is [0.99]
    4. epoch: 2, batch: 0, loss is: [0.0282321], acc is [0.99]
    5. epoch: 2, batch: 200, loss is: [0.04439911], acc is [0.99]
    6. epoch: 2, batch: 400, loss is: [0.03719958], acc is [0.98]
    7. epoch: 3, batch: 0, loss is: [0.04289422], acc is [0.99]
    8. epoch: 3, batch: 200, loss is: [0.02114591], acc is [0.99]
    9. epoch: 3, batch: 400, loss is: [0.00777181], acc is [1.]
    10. epoch: 4, batch: 0, loss is: [0.01535593], acc is [0.99]
    11. epoch: 4, batch: 200, loss is: [0.00227343], acc is [1.]

    从恢复训练的损失变化来看,加载模型参数继续训练的损失函数值和正常训练损失函数值是完全一致的,可见使用飞桨实现恢复训练是极其简单的。

    epoch: 1, batch: 0, loss is: [0.04537241], acc is [0.98] epoch: 1, batch: 200, loss is: [0.04753707], acc is [0.98] epoch: 1, batch: 400, loss is: [0.02443215], acc is [0.99]