近期做了一些反垃圾的工作，除了使用常用的規(guī)則匹配過濾等手段，也采用了一些機(jī)器學(xué)習(xí)方法進(jìn)行分類預(yù)測。我們使用TensorFlow進(jìn)行模型的訓(xùn)練，訓(xùn)練好的模型需要保存，預(yù)測階段我們需要將模型進(jìn)行加載還原使用，這就涉及TensorFlow模型的保存與恢復(fù)加載。

總結(jié)一下Tensorflow常用的模型保存方式。

保存checkpoint模型文件（.ckpt）

首先，TensorFlow提供了一個(gè)非常方便的api，tf.train.Saver()來保存和還原一個(gè)機(jī)器學(xué)習(xí)模型。

模型保存

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

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))

程序生成并保存四個(gè)文件（在版本0.11之前只會(huì)生成三個(gè)文件：checkpoint, model.ckpt, model.ckpt.meta）

1. checkpoint 文本文件，記錄了模型文件的路徑信息列表
2. model.ckpt.data-00000-of-00001 網(wǎng)絡(luò)權(quán)重信息
3. model.ckpt.index .data和.index這兩個(gè)文件是二進(jìn)制文件，保存了模型中的變量參數(shù)（權(quán)重）信息
4. model.ckpt.meta 二進(jìn)制文件，保存了模型的計(jì)算圖結(jié)構(gòu)信息（模型的網(wǎng)絡(luò)結(jié)構(gòu)）protobuf

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

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

加上global_step參數(shù)代表在每1000次迭代后保存模型，會(huì)在模型文件后加上"-1000"，model.ckpt-1000.index, model.ckpt-1000.meta, model.ckpt.data-1000-00000-of-00001

每1000次迭代保存一次模型，但是模型的結(jié)構(gòu)信息文件不會(huì)變，就只用1000次迭代時(shí)保存一下，不用相應(yīng)的每1000次保存一次，所以當(dāng)我們不需要保存meta文件時(shí)，可以加上write_meta_graph=False參數(shù)，如下：

如果想每兩小時(shí)保存一次模型，并且只保存最新的4個(gè)模型，可以加上使用max_to_keep（默認(rèn)值為5，如果想每訓(xùn)練一個(gè)epoch就保存一次，可以將其設(shè)置為None或0，但是沒啥用不推薦）, keep_checkpoint_every_n_hours參數(shù)，如下：

同時(shí)在tf.train.Saver()類中，如果我們不指定任何信息，則會(huì)保存所有的參數(shù)信息，我們也可以指定部分想要保存的內(nèi)容，例如只保存x, y參數(shù)（可傳入?yún)?shù)list或dict）：

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

ps. 在模型訓(xùn)練過程中需要在保存后拿到的變量或參數(shù)名屬性name不能丟，不然模型還原后不能通過get_tensor_by_name()獲取。

模型加載還原

針對(duì)上面的模型保存例子，還原模型的過程如下：

import tensorflow as tf

def restore_model_ckpt(ckpt_file_path):
  sess = tf.Session()
  saver = tf.train.import_meta_graph('./ckpt/model.ckpt.meta') # 加載模型結(jié)構(gòu)
  saver.restore(sess, tf.train.latest_checkpoint('./ckpt')) # 只需要指定目錄就可以恢復(fù)所有變量信息

  # 直接獲取保存的變量
  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')
  # 獲取需要進(jìn)行計(jì)算的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)

首先還原模型結(jié)構(gòu)，然后還原變量（參數(shù)）信息，最后我們就可以獲得已訓(xùn)練的模型中的各種信息了（保存的變量、placeholder變量、operator等），同時(shí)可以對(duì)獲取的變量添加各種新的操作（見以上代碼注釋）。
并且，我們也可以加載部分模型，在此基礎(chǔ)上加入其它操作，具體可以參考官方文檔和demo。

針對(duì)ckpt模型文件的保存與還原，stackoverflow上有一個(gè)回答解釋比較清晰，可以參考。

同時(shí)cv-tricks.com上面的TensorFlow模型保存與恢復(fù)的教程也非常好，可以參考。

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

保存單個(gè)模型文件（.pb）

我自己運(yùn)行過Tensorflow的inception-v3的demo，發(fā)現(xiàn)運(yùn)行結(jié)束后會(huì)生成一個(gè).pb的模型文件，這個(gè)文件是作為后續(xù)預(yù)測或遷移學(xué)習(xí)使用的，就一個(gè)文件，非常炫酷，也十分方便。

這個(gè)過程的主要思路是graph_def文件中沒有包含網(wǎng)絡(luò)中的Variable值（通常情況存儲(chǔ)了權(quán)重），但是卻包含了constant值，所以如果我們能把Variable轉(zhuǎn)換為constant（使用graph_util.convert_variables_to_constants()函數(shù)），即可達(dá)到使用一個(gè)文件同時(shí)存儲(chǔ)網(wǎng)絡(luò)架構(gòu)與權(quán)重的目標(biāo)。

ps：這里.pb是模型文件的后綴名，當(dāng)然我們也可以用其它的后綴（使用.pb與google保持一致 ╮(╯▽╰)╭）

模型保存

同樣根據(jù)上面的例子，一個(gè)簡單的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()，可以多個(gè)
  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))

程序生成并保存一個(gè)文件

model.pb 二進(jìn)制文件，同時(shí)保存了模型網(wǎng)絡(luò)結(jié)構(gòu)和參數(shù)（權(quán)重）信息

模型加載還原

針對(duì)上面的模型保存例子，還原模型的過程如下：

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的網(wǎng)絡(luò)導(dǎo)出為單個(gè)文件》上介紹了TensorFlow保存單個(gè)模型文件的方式，大同小異，可以看看。

思考

模型的保存與加載只是TensorFlow中最基礎(chǔ)的部分之一，雖然簡單但是也必不可少，在實(shí)際運(yùn)用中還需要注意模型何時(shí)保存，哪些變量需要保存，如何設(shè)計(jì)加載實(shí)現(xiàn)遷移學(xué)習(xí)等等問題。

同時(shí)TensorFlow的函數(shù)和類都在一直變化更新，以后也有可能出現(xiàn)更豐富的模型保存和還原的方法。

選擇保存為checkpoint或單個(gè)pb文件視業(yè)務(wù)情況而定，沒有特別大的差別。checkpoint保存感覺會(huì)更加靈活一些，pb文件更適合線上部署吧（個(gè)人看法）。

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

以上就是本文的全部內(nèi)容，希望對(duì)大家的學(xué)習(xí)有所幫助，也希望大家多多支持腳本之家。