我們使用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）
- checkpoint 文本文件，記錄了模型文件的路徑信息列表
- model.ckpt.data-00000-of-00001 網絡權重信息
- model.ckpt.index .data和.index這兩個文件是二進制文件，保存了模型中的變量參數（權重）信息
- 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參數，如下：
如果想每兩小時保存一次模型，並且只保存最新的4個模型，可以加上使用max_to_keep（默認值為5，如果想每訓練一個epoch就保存一次，可以將其設置為None或0，但是沒啥用不推薦）, keep_checkpoint_every_n_hours參數，如下：
同時在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