Tensorflow模型的保存與恢復加載

Blog的Github地址：github.com/liuyan731/b…

---

近期作了一些反垃圾的工做，除了使用經常使用的規則匹配過濾等手段，也採用了一些機器學習方法進行分類預測。咱們使用TensorFlow進行模型的訓練，訓練好的模型須要保存，預測階段咱們須要將模型進行加載還原使用，這就涉及TensorFlow模型的保存與恢復加載。

總結一下Tensorflow經常使用的模型保存方式。

保存checkpoint模型文件（.ckpt）

首先，TensorFlow提供了一個很是方便的api，tf.train.Saver()來保存和還原一個機器學習模型。

模型保存

使用tf.train.Saver()來保存模型文件很是方便，下面是一個簡單的例子：

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()方法還有不少可配置的參數：

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）：

tf.train.Saver([x, y]).save(sess, ckpt_file_path)
複製代碼

ps. 在模型訓練過程當中須要在保存後拿到的變量或參數名屬性name不能丟，否則模型還原後不能經過get_tensor_by_name()獲取。

模型加載還原

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

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模型保存與恢復的教程也很是好，能夠參考。

《tensorflow 1.0 學習：模型的保存與恢復(Saver)》有一些Saver使用技巧。

保存單個模型文件（.pb）

我本身運行過Tensorflow的inception-v3的demo，發現運行結束後會生成一個.pb的模型文件，這個文件是做爲後續預測或遷移學習使用的，就一個文件，很是炫酷，也十分方便。

這個過程的主要思路是graph_def文件中沒有包含網絡中的Variable值（一般狀況存儲了權重），可是卻包含了constant值，因此若是咱們能把Variable轉換爲constant（使用graph_util.convert_variables_to_constants()函數），便可達到使用一個文件同時存儲網絡架構與權重的目標。

ps：這裏.pb是模型文件的後綴名，固然咱們也能夠用其它的後綴（使用.pb與google保持一致 ╮(╯▽╰)╭）

模型保存

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

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 二進制文件，同時保存了模型網絡結構和參數（權重）信息

模型加載還原

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

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差很少同樣。

CSDN《將TensorFlow的網絡導出爲單個文件》上介紹了TensorFlow保存單個模型文件的方式，大同小異，能夠看看。

思考

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

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

選擇保存爲checkpoint或單個pb文件視業務狀況而定，沒有特別大的差異。checkpoint保存感受會更加靈活一些，pb文件更適合線上部署吧（我的見解）。

以上完整代碼：github

---

2017/11/25 done

此文章也同步至我的Github博客