1. 程式人生 > >TensorFlow儲存和提取模型

TensorFlow儲存和提取模型

一、模型的儲存

將訓練好的模型引數儲存起來,以便以後進行驗證或測試,這是我們經常要做的事情。tf裡面提供模型儲存的是tf.train.Saver()模組。
1、模型儲存,先要建立一個Saver物件:如

saver=tf.train.Saver()
__init__(
    var_list=None,
    reshape=False,
    sharded=False,
    max_to_keep=5,
    keep_checkpoint_every_n_hours=10000.0,
    name=None,
    restore_sequentially=False,
    saver_def=None,
    builder=None,
    defer_build=False,
    allow_empty=False,
    write_version=tf.train.SaverDef.V2,
    pad_step_number=False,
    save_relative_paths=False,
    filename=None
)

在建立這個Saver物件的時候,有一個引數我們經常會用到,就是 max_to_keep 引數,這個是用來設定儲存模型的個數,預設為5,即 max_to_keep=5,儲存最近的5個模型。如果你想每訓練一代(epoch)就想儲存一次模型,則可以將 max_to_keep設定為None或者0,如:

saver=tf.train.Saver(max_to_keep=0)

但是這樣做除了多佔用硬碟,並沒有實際多大的用處,因此不推薦。
當然,如果你只想儲存最後一代的模型,則只需要將max_to_keep設定為1即可,即

saver=tf.train.Saver(max_to_keep=1)

某個程式的例子:

model_saver=tf.train.Saver(var_list=TEM_trainable_variables,max_to_keep=80)

2、建立完saver物件後,就可以儲存訓練好的模型了,如:

saver.save(sess,'ckpt/mnist.ckpt',global_step=step)

第一個引數sess,這個就不用說了。第二個引數設定儲存的路徑和名字,第三個引數將訓練的次數作為字尾加入到模型名字中。

saver.save(sess, 'my-model', global_step=0) ==> filename: 'my-model-0'
...
saver.save(sess, 'my-model', global_step=1000) ==> filename: 'my-model-1000'

幾個程式的例子:

model_saver.save(sess,"models/TEM/tem_model_checkpoint")
saver=tf.train.Saver(max_to_keep=1)
for i in range(100):
  batch_xs, batch_ys = mnist.train.next_batch(100)
  sess.run(train_op, feed_dict={x: batch_xs, y_: batch_ys})
  val_loss,val_acc=sess.run([loss,acc], feed_dict={x: mnist.test.images, y_: mnist.test.labels})
  print('epoch:%d, val_loss:%f, val_acc:%f'%(i,val_loss,val_acc))
  saver.save(sess,'ckpt/mnist.ckpt',global_step=i+1)
sess.close()

程式碼中紅色部分就是儲存模型的程式碼,雖然我在每訓練完一代的時候,都進行了儲存,但後一次儲存的模型會覆蓋前一次的,最終只會儲存最後一次。因此我們可以節省時間,將儲存程式碼放到迴圈之外(僅適用max_to_keep=1,否則還是需要放在迴圈內).

在實驗中,最後一代可能並不是驗證精度最高的一代,因此我們並不想預設儲存最後一代,而是想儲存驗證精度最高的一代,則加個中間變數和判斷語句就可以了。

saver=tf.train.Saver(max_to_keep=1)
max_acc=0
for i in range(100):
  batch_xs, batch_ys = mnist.train.next_batch(100)
  sess.run(train_op, feed_dict={x: batch_xs, y_: batch_ys})
  val_loss,val_acc=sess.run([loss,acc], feed_dict={x: mnist.test.images, y_: mnist.test.labels})
  print('epoch:%d, val_loss:%f, val_acc:%f'%(i,val_loss,val_acc))
  if val_acc>max_acc:
      max_acc=val_acc
      saver.save(sess,'ckpt/mnist.ckpt',global_step=i+1)
sess.close()

如果我們想儲存驗證精度最高的三代,且把每次的驗證精度也隨之儲存下來,則我們可以生成一個txt檔案用於儲存。

saver=tf.train.Saver(max_to_keep=3)
max_acc=0
f=open('ckpt/acc.txt','w')
for i in range(100):
  batch_xs, batch_ys = mnist.train.next_batch(100)
  sess.run(train_op, feed_dict={x: batch_xs, y_: batch_ys})
  val_loss,val_acc=sess.run([loss,acc], feed_dict={x: mnist.test.images, y_: mnist.test.labels})
  print('epoch:%d, val_loss:%f, val_acc:%f'%(i,val_loss,val_acc))
  f.write(str(i+1)+', val_acc: '+str(val_acc)+'\n')
  if val_acc>max_acc:
      max_acc=val_acc
      saver.save(sess,'ckpt/mnist.ckpt',global_step=i+1)
f.close()
sess.close()

二、模型的提取

模型的恢復用的是restore()函式,它需要兩個引數restore(sess, save_path),save_path指的是儲存的模型路徑。我們可以使用tf.train.latest_checkpoint()來自動獲取最後一次儲存的模型。如:

model_file=tf.train.latest_checkpoint('ckpt/')
saver.restore(sess,model_file)

某個程式的例子:

model_saver.restore(sess,"models/TEM/tem_model_best")