近期做了一些反垃圾的工作,除了使用常用的規則匹配過濾等手段,也採用了一些機器學習方法進行分類預測。我們使用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模型儲存與恢復的教程也非常好,可以參考。

儲存單個模型檔案(.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檔案更適合線上部署吧(個人看法)。