TensorFlow——Checkpoint為模型新增檢查點

Lily發表於2019-05-27
模型

1.檢查點

儲存模型並不限於在訓練模型後，在訓練模型之中也需要儲存，因為TensorFlow訓練模型時難免會出現中斷的情況，我們自然希望能夠將訓練得到的引數儲存下來，否則下次又要重新訓練。

這種在訓練中儲存模型，習慣上稱之為儲存檢查點。

2.新增儲存點

通過新增檢查點，可以生成載入檢查點檔案，並能夠指定生成檢查檔案的個數，例如使用saver的另一個引數——max_to_keep=1，表明最多隻儲存一個檢查點檔案，在儲存時使用如下的程式碼傳入迭代次數。

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import os

train_x = np.linspace(-5, 3, 50)
train_y = train_x * 5 + 10 + np.random.random(50) * 10 - 5

plt.plot(train_x, train_y, 'r.')
plt.grid(True)
plt.show()

tf.reset_default_graph()

X = tf.placeholder(dtype=tf.float32)
Y = tf.placeholder(dtype=tf.float32)

w = tf.Variable(tf.random.truncated_normal([1]), name='Weight')
b = tf.Variable(tf.random.truncated_normal([1]), name='bias')

z = tf.multiply(X, w) + b

cost = tf.reduce_mean(tf.square(Y - z))
learning_rate = 0.01
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost)

init = tf.global_variables_initializer()

training_epochs = 20
display_step = 2


saver = tf.train.Saver(max_to_keep=15)
savedir = "model/"


if __name__ == '__main__':
    with tf.Session() as sess:
        sess.run(init)
        loss_list = []
        for epoch in range(training_epochs):
            for (x, y) in zip(train_x, train_y):
                sess.run(optimizer, feed_dict={X: x, Y: y})

            if epoch % display_step == 0:
                loss = sess.run(cost, feed_dict={X: x, Y: y})
                loss_list.append(loss)
                print('Iter: ', epoch, ' Loss: ', loss)

            w_, b_ = sess.run([w, b], feed_dict={X: x, Y: y})

            saver.save(sess, savedir + "linear.cpkt", global_step=epoch)

        print(" Finished ")
        print("W: ", w_, " b: ", b_, " loss: ", loss)
        plt.plot(train_x, train_x * w_ + b_, 'g-', train_x, train_y, 'r.')
        plt.grid(True)
        plt.show()

    load_epoch = 10

    with tf.Session() as sess2:
        sess2.run(tf.global_variables_initializer())
        saver.restore(sess2, savedir + "linear.cpkt-" + str(load_epoch))
        print(sess2.run([w, b], feed_dict={X: train_x, Y: train_y}))

在上述的程式碼中，我們使用saver.save(sess, savedir + "linear.cpkt", global_step=epoch)將訓練的引數傳入檢查點進行儲存，saver = tf.train.Saver(max_to_keep=1)表示只儲存一個檔案，這樣在訓練過程中得到的新的模型就會覆蓋以前的模型。

cpkt = tf.train.get_checkpoint_state(savedir)
if cpkt and cpkt.model_checkpoint_path:
　　saver.restore(sess2, cpkt.model_checkpoint_path)

kpt = tf.train.latest_checkpoint(savedir)
saver.restore(sess2, kpt)

上述的兩種方法也可以對checkpoint檔案進行載入，tf.train.latest_checkpoint(savedir)為載入最後的檢查點檔案。這種方式，我們可以通過儲存指定訓練次數的檢查點，比如儲存5的倍數次儲存一下檢查點。

3.簡便儲存檢查點

我們還可以用更加簡單的方法進行檢查點的儲存，tf.train.MonitoredTrainingSession()函式，該函式可以直接實現儲存載入檢查點模型的檔案，與前面的方法不同的是，它是按照訓練時間來儲存檢查點的，可以通過指定save_checkpoint_secs引數的具體秒數，設定多久儲存一次檢查點。

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import os

train_x = np.linspace(-5, 3, 50)
train_y = train_x * 5 + 10 + np.random.random(50) * 10 - 5

# plt.plot(train_x, train_y, 'r.')
# plt.grid(True)
# plt.show()

tf.reset_default_graph()

X = tf.placeholder(dtype=tf.float32)
Y = tf.placeholder(dtype=tf.float32)

w = tf.Variable(tf.random.truncated_normal([1]), name='Weight')
b = tf.Variable(tf.random.truncated_normal([1]), name='bias')

z = tf.multiply(X, w) + b

cost = tf.reduce_mean(tf.square(Y - z))
learning_rate = 0.01
optimizer = tf.train.GradientDescentOptimizer(learning_rate).minimize(cost)

init = tf.global_variables_initializer()

training_epochs = 30
display_step = 2


global_step = tf.train.get_or_create_global_step()

step = tf.assign_add(global_step, 1)

saver = tf.train.Saver()

savedir = "check-point/"

if __name__ == '__main__':
    with tf.train.MonitoredTrainingSession(checkpoint_dir=savedir + 'linear.cpkt', save_checkpoint_secs=5) as sess:
        sess.run(init)
        loss_list = []
        for epoch in range(training_epochs):
            sess.run(global_step)
            for (x, y) in zip(train_x, train_y):
                sess.run(optimizer, feed_dict={X: x, Y: y})

            if epoch % display_step == 0:
                loss = sess.run(cost, feed_dict={X: x, Y: y})
                loss_list.append(loss)
                print('Iter: ', epoch, ' Loss: ', loss)

            w_, b_ = sess.run([w, b], feed_dict={X: x, Y: y})
            sess.run(step)

        print(" Finished ")
        print("W: ", w_, " b: ", b_, " loss: ", loss)
        plt.plot(train_x, train_x * w_ + b_, 'g-', train_x, train_y, 'r.')
        plt.grid(True)
        plt.show()

    load_epoch = 10

    with tf.Session() as sess2:
        sess2.run(tf.global_variables_initializer())

        # saver.restore(sess2, savedir + 'linear.cpkt-' + str(load_epoch))

        # cpkt = tf.train.get_checkpoint_state(savedir)
        # if cpkt and cpkt.model_checkpoint_path:
        #     saver.restore(sess2, cpkt.model_checkpoint_path)
        #
        kpt = tf.train.latest_checkpoint(savedir + 'linear.cpkt')

        saver.restore(sess2, kpt)

        print(sess2.run([w, b], feed_dict={X: train_x, Y: train_y}))

上述的程式碼中，我們設定了沒訓練了5秒中之後，就儲存一次檢查點，它預設的儲存時間間隔是10分鐘，這種按照時間的儲存模式更適合使用大型資料集訓練複雜模型的情況，注意在使用上述的方法時，要定義global_step變數，在訓練完一個批次或者一個樣本之後，要將其進行加1的操作，否則將會報錯。

 

相關文章