將訓練好的模型保存到磁盤之后,應用程序可以隨時加載模型,完成預測任務。但是在日常訓練工作中我們會遇到一些突發情況,導致訓練過程主動或被動的中斷。如果訓練一個模型需要花費幾天的訓練時間,中斷后從初始狀態重新訓練是不可接受的。
飛槳支持從上一次保存狀態開始訓練,只要我們隨時保存訓練過程中的模型狀態,就不用從初始狀態重新訓練。
下面介紹恢復訓練的代碼實現,依然使用手寫數字識別的案例,在網絡定義的部分保持不變。

1 import os 2 import random 3 import paddle 4 import paddle.fluid as fluid 5 from paddle.fluid.dygraph.nn import Conv2D, Pool2D, FC 6 import numpy as np 7 from PIL import Image 8 9 import gzip 10 import json 11 12 # 定義數據集讀取器 13 def load_data(mode='train'): 14 15 # 數據文件 16 datafile = './work/mnist.json.gz' 17 print('loading mnist dataset from {} ......'.format(datafile)) 18 data = json.load(gzip.open(datafile)) 19 train_set, val_set, eval_set = data 20 21 # 數據集相關參數,圖片高度IMG_ROWS, 圖片寬度IMG_COLS 22 IMG_ROWS = 28 23 IMG_COLS = 28 24 25 if mode == 'train': 26 imgs = train_set[0] 27 labels = train_set[1] 28 elif mode == 'valid': 29 imgs = val_set[0] 30 labels = val_set[1] 31 elif mode == 'eval': 32 imgs = eval_set[0] 33 labels = eval_set[1] 34 35 imgs_length = len(imgs) 36 37 assert len(imgs) == len(labels), \ 38 "length of train_imgs({}) should be the same as train_labels({})".format( 39 len(imgs), len(labels)) 40 41 index_list = list(range(imgs_length)) 42 43 # 讀入數據時用到的batchsize 44 BATCHSIZE = 100 45 46 # 定義數據生成器 47 def data_generator(): 48 #if mode == 'train': 49 # random.shuffle(index_list) 50 imgs_list = [] 51 labels_list = [] 52 for i in index_list: 53 img = np.reshape(imgs[i], [1, IMG_ROWS, IMG_COLS]).astype('float32') 54 label = np.reshape(labels[i], [1]).astype('int64') 55 imgs_list.append(img) 56 labels_list.append(label) 57 if len(imgs_list) == BATCHSIZE: 58 yield np.array(imgs_list), np.array(labels_list) 59 imgs_list = [] 60 labels_list = [] 61 62 # 如果剩余數據的數目小於BATCHSIZE, 63 # 則剩余數據一起構成一個大小為len(imgs_list)的mini-batch 64 if len(imgs_list) > 0: 65 yield np.array(imgs_list), np.array(labels_list) 66 67 return data_generator 68 69 #調用加載數據的函數 70 train_loader = load_data('train') 71 72 # 定義模型結構 73 class MNIST(fluid.dygraph.Layer): 74 def __init__(self, name_scope): 75 super(MNIST, self).__init__(name_scope) 76 name_scope = self.full_name() 77 self.conv1 = Conv2D(name_scope, num_filters=20, filter_size=5, stride=1, padding=2, act="relu") 78 self.pool1 = Pool2D(name_scope, pool_size=2, pool_stride=2, pool_type='max') 79 self.conv2 = Conv2D(name_scope, num_filters=20, filter_size=5, stride=1, padding=2, act="relu") 80 self.pool2 = Pool2D(name_scope, pool_size=2, pool_stride=2, pool_type='max') 81 self.fc = FC(name_scope, size=10, act='softmax') 82 83 #加入分類准確率的評估指標 84 def forward(self, inputs, label=None): 85 x = self.conv1(inputs) 86 x = self.pool1(x) 87 x = self.conv2(x) 88 x = self.pool2(x) 89 x = self.fc(x) 90 if label is not None: 91 acc = fluid.layers.accuracy(input=x, label=label) 92 return x, acc 93 else: 94 return x
loading mnist dataset from ./work/mnist.json.gz ......
這是上面打印輸出的一句話。
在開始使用飛槳恢復訓練前,先正常訓練一個模型,優化器使用Adam,使用動態變化的學習率,學習率從0.01衰減到0.001(12-16行)。每訓練一輪后保存一次模型,之后將采用訓練中的模型參數作為恢復訓練的模型參數繼續訓練。
說明:
本次訓練不僅保存模型參數,而且保存優化器、學習率有關的參數,比如Adam, Adagrad優化器在訓練時會創建一些新的變量輔助訓練;動態變化的學習率需要訓練停止時的訓練步數。這些參數對於恢復訓練至關重要。
保存模型參數見43行,保存優化器參數見44行。
1 #在使用GPU機器時,可以將use_gpu變量設置成True 2 use_gpu = True 3 place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace() 4 5 with fluid.dygraph.guard(place): 6 model = MNIST("mnist") 7 model.train() 8 9 EPOCH_NUM = 5 10 BATCH_SIZE = 100 11 # 定義學習率,並加載優化器參數到模型中 12 total_steps = (int(60000//BATCH_SIZE) + 1) * EPOCH_NUM 13 lr = fluid.dygraph.PolynomialDecay(0.01, total_steps, 0.001) 14 15 # 使用Adam優化器 16 optimizer = fluid.optimizer.AdamOptimizer(learning_rate=lr) 17 18 for epoch_id in range(EPOCH_NUM): 19 for batch_id, data in enumerate(train_loader()): 20 #准備數據,變得更加簡潔 21 image_data, label_data = data 22 image = fluid.dygraph.to_variable(image_data) 23 label = fluid.dygraph.to_variable(label_data) 24 25 #前向計算的過程,同時拿到模型輸出值和分類准確率 26 predict, acc = model(image, label) 27 avg_acc = fluid.layers.mean(acc) 28 29 #計算損失,取一個批次樣本損失的平均值 30 loss = fluid.layers.cross_entropy(predict, label) 31 avg_loss = fluid.layers.mean(loss) 32 33 #每訓練了200批次的數據,打印下當前Loss的情況 34 if batch_id % 200 == 0: 35 print("epoch: {}, batch: {}, loss is: {}, acc is {}".format(epoch_id, batch_id, avg_loss.numpy(),avg_acc.numpy())) 36 37 #后向傳播,更新參數的過程 38 avg_loss.backward() 39 optimizer.minimize(avg_loss) 40 model.clear_gradients() 41 42 # 保存模型參數和優化器的參數 43 fluid.save_dygraph(model.state_dict(), './checkpoint/mnist_epoch{}'.format(epoch_id)) 44 fluid.save_dygraph(optimizer.state_dict(), './checkpoint/mnist_epoch{}'.format(epoch_id))
epoch: 0, batch: 0, loss is: [2.8080773], acc is [0.16] epoch: 0, batch: 200, loss is: [0.12834078], acc is [0.95] epoch: 0, batch: 400, loss is: [0.10862964], acc is [0.97] epoch: 1, batch: 0, loss is: [0.09647215], acc is [0.98] epoch: 1, batch: 200, loss is: [0.06849223], acc is [0.97] epoch: 1, batch: 400, loss is: [0.06437591], acc is [0.98] epoch: 2, batch: 0, loss is: [0.0582899], acc is [0.99] epoch: 2, batch: 200, loss is: [0.03881769], acc is [0.99] epoch: 2, batch: 400, loss is: [0.02290947], acc is [0.99] epoch: 3, batch: 0, loss is: [0.08704039], acc is [0.98] epoch: 3, batch: 200, loss is: [0.04802512], acc is [0.99] epoch: 3, batch: 400, loss is: [0.01132718], acc is [1.] epoch: 4, batch: 0, loss is: [0.07382318], acc is [0.99] epoch: 4, batch: 200, loss is: [0.05015907], acc is [0.98] epoch: 4, batch: 400, loss is: [0.0086437], acc is [1.]
恢復訓練
在上述訓練代碼中,我們訓練了五輪(epoch)。在每輪結束時,我們均保存了模型參數和優化器相關的參數。
- 使用model.state_dict()獲取模型參數。
- 使用optimizer.state_dict()獲取優化器和學習率相關的參數。
- 調用paddle的save_dygraph API將參數保存到本地。
比如第一輪訓練保存的文件是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 params_path = "./checkpoint/mnist_epoch2" 2 #在使用GPU機器時,可以將use_gpu變量設置成True 3 use_gpu = True 4 place = fluid.CUDAPlace(0) if use_gpu else fluid.CPUPlace() 5 6 with fluid.dygraph.guard(place): 7 # 加載模型參數到模型中 8 params_dict, opt_dict = fluid.load_dygraph(params_path) 9 model = MNIST("mnist") 10 model.load_dict(params_dict) 11 12 EPOCH_NUM = 5 13 BATCH_SIZE = 100 14 # 定義學習率,並加載優化器參數到模型中 15 total_steps = (int(60000//BATCH_SIZE) + 1) * EPOCH_NUM 16 lr = fluid.dygraph.PolynomialDecay(0.01, total_steps, 0.001) 17 18 # 使用Adam優化器 19 optimizer = fluid.optimizer.AdamOptimizer(learning_rate=lr) 20 optimizer.set_dict(opt_dict) 21 22 for epoch_id in range(3, EPOCH_NUM): 23 for batch_id, data in enumerate(train_loader()): 24 #准備數據,變得更加簡潔 25 image_data, label_data = data 26 image = fluid.dygraph.to_variable(image_data) 27 label = fluid.dygraph.to_variable(label_data) 28 29 #前向計算的過程,同時拿到模型輸出值和分類准確率 30 predict, acc = model(image, label) 31 avg_acc = fluid.layers.mean(acc) 32 33 #計算損失,取一個批次樣本損失的平均值 34 loss = fluid.layers.cross_entropy(predict, label) 35 avg_loss = fluid.layers.mean(loss) 36 37 #每訓練了200批次的數據,打印下當前Loss的情況 38 if batch_id % 200 == 0: 39 print("epoch: {}, batch: {}, loss is: {}, acc is {}".format(epoch_id, batch_id, avg_loss.numpy(),avg_acc.numpy())) 40 41 #后向傳播,更新參數的過程 42 avg_loss.backward() 43 optimizer.minimize(avg_loss) 44 model.clear_gradients()
epoch: 3, batch: 0, loss is: [0.08704039], acc is [0.98] epoch: 3, batch: 200, loss is: [0.03669801], acc is [0.99] epoch: 3, batch: 400, loss is: [0.0107842], acc is [1.] epoch: 4, batch: 0, loss is: [0.06170184], acc is [0.99] epoch: 4, batch: 200, loss is: [0.0378595], acc is [0.99] epoch: 4, batch: 400, loss is: [0.01447322], acc is [1.]
我更改了第1行和第22行,使系統從第3個epoch開始訓練。
如果從第一個epoch開始重新訓練,則1和22行改為*_epoch0,1這兩個地方。
從恢復訓練的損失變化來看,加載模型參數繼續訓練的損失函數值和正常訓練損失函數值是完全一致的,可見使用飛槳實現恢復訓練是極其簡單的。