TensorFlow模型的儲存與恢複加載 ckpt

我們使用TensorFlow進行模型的訓練，訓練好的模型需要儲存，預測階段我們需要将模型進行加載還原使用，這就涉及TensorFlow模型的儲存與恢複加載。

總結一下Tensorflow常用的模型儲存方式。

儲存checkpoint模型檔案（.ckpt）

首先，TensorFlow提供了一個非常友善的api，tf.train.Saver()來儲存和還原一個機器學習模型。

舊版儲存單個模型檔案的方式，大同小異，可以看看。

• .ckpt檔案是舊版本的輸出

  saver.save(sess)

  ，相當于你的

  .ckpt-data

  （見下文）
• “checkpoint”檔案僅用于告知某些TF函數，這是最新的檢查點檔案。

新版儲存為：

• • .ckpt.meta
  • .ckpt.index
  • .ckpt.data

• .ckpt-meta

   包含元圖，即計算圖的結構，沒有變量的值（基本上你可以在tensorboard / graph中看到）。

• .ckpt-data

  包含所有變量的值，沒有結構。要在python中恢複模型，您通常會使用中繼資料和資料檔案（但您也可以使用該

  .pb

  檔案）： 

  saver = tf.train.import_meta_graph(path_to_ckpt_meta) saver.restore(sess, path_to_ckpt_data)

• 我不确切地知道

  .ckpt-index

  ，我想這是内部需要的某種索引來正确映射前兩個檔案。
• 該

  .pb

  檔案可以儲存您的整個圖表（元+資料）。要在c ++中加載和使用（但不訓練）圖形，您通常會使用它來建立

  freeze_graph

  ，它會

  .pb

  從中繼資料和資料建立檔案。要小心，（至少在以前的TF版本和某些人中）py提供的功能

  freeze_graph

  不能正常工作，是以你必須使用腳本版本。Tensorflow還提供了一種

  tf.train.Saver.to_proto()

  方法，但我不知道它究竟是做什麼的。

模型儲存

使用tf.train.Saver()來儲存模型檔案非常友善，下面是一個簡單的例子：

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22

import tensorflow as tf import os def save_model_ckpt(ckpt_file_path): x = tf.placeholder(tf.int32, name = 'x' ) y = tf.placeholder(tf.int32, name = 'y' ) b = tf.Variable( 1 , name = 'b' ) xy = tf.multiply(x, y) op = tf.add(xy, b, name = 'op_to_store' ) sess = tf.Session() sess.run(tf.global_variables_initializer()) path = os.path.dirname(os.path.abspath(ckpt_file_path)) if os.path.isdir(path) is False : os.makedirs(path) tf.train.Saver().save(sess, ckpt_file_path) # test feed_dict = {x: 2 , y: 3 } print (sess.run(op, feed_dict))

程式生成并儲存四個檔案（在版本0.11之前隻會生成三個檔案：checkpoint, model.ckpt, model.ckpt.meta）

1. checkpoint 文本檔案，記錄了模型檔案的路徑資訊清單
2. model.ckpt.data-00000-of-00001 網絡權重資訊
3. model.ckpt.index .data和.index這兩個檔案是二進制檔案，儲存了模型中的變量參數（權重）資訊
4. model.ckpt.meta 二進制檔案，儲存了模型的計算圖結構資訊（模型的網絡結構）protobuf

以上是tf.train.Saver().save()的基本用法，save()方法還有很多可配置的參數：

1	tf.train.Saver().save(sess, ckpt_file_path, global_step = 1000 )

加上global_step參數代表在每1000次疊代後儲存模型，會在模型檔案後加上"-1000"，model.ckpt-1000.index, model.ckpt-1000.meta, model.ckpt.data-1000-00000-of-00001

每1000次疊代儲存一次模型，但是模型的結構資訊檔案不會變，就隻用1000次疊代時儲存一下，不用相應的每1000次儲存一次，是以當我們不需要儲存meta檔案時，可以加上write_meta_graph=False參數，如下：

複制代碼 代碼如下:

tf.train.Saver().save(sess, ckpt_file_path, global_step=1000, write_meta_graph=False)

如果想每兩小時儲存一次模型，并且隻儲存最新的4個模型，可以加上使用max_to_keep（預設值為5，如果想每訓練一個epoch就儲存一次，可以将其設定為None或0，但是沒啥用不推薦）, keep_checkpoint_every_n_hours參數，如下：

複制代碼 代碼如下:

tf.train.Saver().save(sess, ckpt_file_path, max_to_keep=4, keep_checkpoint_every_n_hours=2)

同時在tf.train.Saver()類中，如果我們不指定任何資訊，則會儲存所有的參數資訊，我們也可以指定部分想要儲存的内容，例如隻儲存x, y參數（可傳入參數list或dict）：

1	tf.train.Saver([x, y]).save(sess, ckpt_file_path)

ps. 在模型訓練過程中需要在儲存後拿到的變量或參數名屬性name不能丢，不然模型還原後不能通過get_tensor_by_name()擷取。

模型加載還原

針對上面的模型儲存例子，還原模型的過程如下：

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21

import tensorflow as tf def restore_model_ckpt(ckpt_file_path): sess = tf.Session() saver = tf.train.import_meta_graph( './ckpt/model.ckpt.meta' ) # 加載模型結構 saver.restore(sess, tf.train.latest_checkpoint( './ckpt' )) # 隻需要指定目錄就可以恢複所有變量資訊 # 直接擷取儲存的變量 print (sess.run( 'b:0' )) # 擷取placeholder變量 input_x = sess.graph.get_tensor_by_name( 'x:0' ) input_y = sess.graph.get_tensor_by_name( 'y:0' ) # 擷取需要進行計算的operator op = sess.graph.get_tensor_by_name( 'op_to_store:0' ) # 加入新的操作 add_on_op = tf.multiply(op, 2 ) ret = sess.run(add_on_op, {input_x: 5 , input_y: 5 }) print (ret)

首先還原模型結構，然後還原變量（參數）資訊，最後我們就可以獲得已訓練的模型中的各種資訊了（儲存的變量、placeholder變量、operator等），同時可以對擷取的變量添加各種新的操作（見以上代碼注釋）。

并且，我們也可以加載部分模型，在此基礎上加入其它操作，具體可以參考官方文檔和demo。

針對ckpt模型檔案的儲存與還原，stackoverflow上有一個回答解釋比較清晰，可以參考。

同時cv-tricks.com上面的TensorFlow模型儲存與恢複的教程也非常好，可以參考。https://stackoverflow.com/questions/33759623/tensorflow-how-to-save-restore-a-model

《tensorflow 1.0 學習：模型的儲存與恢複(Saver)》有一些Saver使用技巧。https://www.jb51.net/article/138781.html

儲存單個模型檔案（.pb）

我自己運作過Tensorflow的inception-v3的demo，發現運作結束後會生成一個.pb的模型檔案，這個檔案是作為後續預測或遷移學習使用的，就一個檔案，非常炫酷，也十分友善。

這個過程的主要思路是graph_def檔案中沒有包含網絡中的Variable值（通常情況存儲了權重），但是卻包含了constant值，是以如果我們能把Variable轉換為constant（使用graph_util.convert_variables_to_constants()函數），即可達到使用一個檔案同時存儲網絡架構與權重的目标。

ps：這裡.pb是模型檔案的字尾名，當然我們也可以用其它的字尾（使用.pb與google保持一緻 ╮(╯▽╰)╭）

模型儲存

同樣根據上面的例子，一個簡單的demo：

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27

import tensorflow as tf import os from tensorflow.python.framework import graph_util def save_mode_pb(pb_file_path): x = tf.placeholder(tf.int32, name = 'x' ) y = tf.placeholder(tf.int32, name = 'y' ) b = tf.Variable( 1 , name = 'b' ) xy = tf.multiply(x, y) # 這裡的輸出需要加上name屬性 op = tf.add(xy, b, name = 'op_to_store' ) sess = tf.Session() sess.run(tf.global_variables_initializer()) path = os.path.dirname(os.path.abspath(pb_file_path)) if os.path.isdir(path) is False : os.makedirs(path) # convert_variables_to_constants 需要指定output_node_names，list()，可以多個 constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, [ 'op_to_store' ]) with tf.gfile.FastGFile(pb_file_path, mode = 'wb' ) as f: f.write(constant_graph.SerializeToString()) # test feed_dict = {x: 2 , y: 3 } print (sess.run(op, feed_dict))

程式生成并儲存一個檔案

model.pb 二進制檔案，同時儲存了模型網絡結構和參數（權重）資訊

模型加載還原

針對上面的模型儲存例子，還原模型的過程如下：

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20

import tensorflow as tf from tensorflow.python.platform import gfile def restore_mode_pb(pb_file_path): sess = tf.Session() with gfile.FastGFile(pb_file_path, 'rb' ) as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) sess.graph.as_default() tf.import_graph_def(graph_def, name = '') print (sess.run( 'b:0' )) input_x = sess.graph.get_tensor_by_name( 'x:0' ) input_y = sess.graph.get_tensor_by_name( 'y:0' ) op = sess.graph.get_tensor_by_name( 'op_to_store:0' ) ret = sess.run(op, {input_x: 5 , input_y: 5 }) print (ret)

模型的還原過程與checkpoint差不多一樣。

思考

模型的儲存與加載隻是TensorFlow中最基礎的部分之一，雖然簡單但是也必不可少，在實際運用中還需要注意模型何時儲存，哪些變量需要儲存，如何設計加載實作遷移學習等等問題。

同時TensorFlow的函數和類都在一直變化更新，以後也有可能出現更豐富的模型儲存和還原的方法。

選擇儲存為checkpoint或單個pb檔案視業務情況而定，沒有特别大的差别。checkpoint儲存感覺會更加靈活一些，pb檔案更适合線上部署吧（個人看法）。

以上完整代碼：github https://github.com/liuyan731/tf_demo

參考1