TensorFlow - 保存和加载训练过程

TensorFlow - 保存和加载训练过程

flyfish

保存训练过程

以线性回归为例

import numpy as np
import tensorflow as tf

# 线性回归模型 y = Wx + b
x = tf.placeholder(tf.float32, [None, 1])
W = tf.Variable(tf.zeros([1,1]))
b = tf.Variable(tf.zeros([1]))
product = tf.matmul(x,W)
y = product + b
y_ = tf.placeholder(tf.float32, [None, 1])

# 成本函数 最小二乘法
cost = tf.reduce_mean(tf.square(y_-y))


#使用梯度下降训练以降低成本
train_step = tf.train.GradientDescentOptimizer(0.0000001).minimize(cost)

sess = tf.Session()
init = tf.global_variables_initializer()
sess.run(init)
steps = 10

saver = tf.train.Saver()  

for i in range(steps):
  #为y = W.x + b ， W = 2, b = 0模型创建模拟数据
  xs = np.array([[i]])
  ys = np.array([[2*i]])
  # Train
  feed = { x: xs, y_: ys }
  sess.run(train_step, feed_dict=feed)
  print("After %d iteration:" % i)
  print("W: %f" % sess.run(W))
  print("b: %f" % sess.run(b))

  print("cost: %f" % sess.run(cost, feed_dict=feed))

# W 应该接近 2,  b 应该接近 0
  if(i % 9==0):  
      saver.save(sess, 'C:\\model.data')
      print("success")

查看先前打印的数据，在迭代到第9次时保存
After 9 iteration:
W: 0.000114
b: 0.000018
cost: 323.962433

加载先前训练的结果

import tensorflow as tf

W = tf.Variable(tf.zeros([1,1]))
b = tf.Variable(tf.zeros([1]))

saver=tf.train.Saver()  

with tf.Session() as sess: 
    saver.restore(sess, 'C:\\model.data')
    print("W: %f" % sess.run(W))
    print("b: %f" % sess.run(b))

输出
INFO:tensorflow:Restoring parameters from C:\model.data
W: 0.000114
b: 0.000018

部分代码解释

为了得到可以用来后续恢复模型以进一步训练或评估的检查点文件（checkpoint file），我们实例化一个tf.train.Saver。

saver = tf.train.Saver()
在训练循环中，将定期调用saver.save()方法，向训练文件夹中写入包含了当前所有可训练变量值得检查点文件。

saver.save(sess, FLAGS.train_dir, global_step=step)
这样，我们以后就可以使用saver.restore()方法，重载模型的参数，继续训练。

saver.restore(sess, FLAGS.train_dir)