TensorFlow 模型保存與恢復(fù)

????????上一篇文章 TensorFlow 訓(xùn)練 CNN 分類(lèi)器 中說(shuō)明了訓(xùn)練簡(jiǎn)單 CNN 模型的整個(gè)過(guò)程栏渺,并在訓(xùn)練結(jié)束后使用 .save 函數(shù)來(lái)保存訓(xùn)練的結(jié)果曙强,其后通過(guò)使用 tf.train.import_meta_graph.restore 函數(shù)來(lái)導(dǎo)入模型進(jìn)行推斷腕扶。本文承接上文,對(duì)模型保存與恢復(fù)做一個(gè)總結(jié)。

????????總的來(lái)說(shuō)懒鉴,模型在保存和恢復(fù)時(shí)最重要的是留下數(shù)據(jù)接口犬性,方便使用時(shí)傳入數(shù)據(jù)和獲取結(jié)果瞻离。TensorFlow 中常用的模型保存格式為 .ckpt 和 .pb,下面分別進(jìn)行詳細(xì)說(shuō)明仔夺。

一琐脏、ckpt 格式模型保存與恢復(fù)

????????.ckpt 格式保存與恢復(fù)都很簡(jiǎn)單,具體可參考 TensorFlow 訓(xùn)練 CNN 分類(lèi)器缸兔。

1. ckpt 格式模型保存

inputs = tf.placeholder(tf.float32, shape=[None, ···], name='inputs')  <-- 入口
···
prediction = tf.nn.softmax(logits, name='prediction')  <-- 出口(僅作為例子日裙,下同)
···
saver = tf.train.Saver()
···

with tf.Session() as sess:
    ···    <-- 訓(xùn)練過(guò)程
    saver.save(sess, './xxx/xxx.ckpt')  <-- 模型保存

????????如上述代碼所示,假設(shè)你定義了一個(gè) TensorFlow 模型惰蜜,數(shù)據(jù)入口由占位符 inputs 給定昂拂,結(jié)果出口由張量 prediction 給定。通過(guò)語(yǔ)句 saver = tf.train.Saver() 定義了模型保存的一個(gè)實(shí)例對(duì)象 saver抛猖,當(dāng)模型訓(xùn)練結(jié)束之后只需要簡(jiǎn)單的一條語(yǔ)句:

saver.save(sess, path_to_model.ckpt)

就把訓(xùn)練結(jié)果保存到了指定的路徑格侯。

????????以上代碼之所以把變量 inputsprediction 單獨(dú)列出,一方面是因?yàn)樗鼈兪悄P?Graph 的起點(diǎn)和終點(diǎn)(戲稱(chēng)為數(shù)據(jù)入口财著、出口)联四,另一方面的原因是它們被特別的指定了名稱(chēng),因而在模型恢復(fù)時(shí)可以通過(guò)它們的名稱(chēng)而得到 Graph 中對(duì)應(yīng)的節(jié)點(diǎn)撑教。

2. ckpt 格式模型恢復(fù)

????????當(dāng)你需要導(dǎo)入模型進(jìn)行推斷時(shí)朝墩,只需要通過(guò)張量名獲取數(shù)據(jù)入口和出口,然后傳入數(shù)據(jù)即可:

with tf.Session() as sess:
    saver = tf.train.import_meta_graph('./xxx/xxx.ckpt.meta')
    saver.restore(sess, './xxx/xxx.ckpt')

    inputs = tf.get_default_graph().get_tensor_by_name('inputs:0')
    prediction = tf.get_default_graph().get_tensor_by_name('prediction:0')

    pred = sess.run(prediction, feed_dict={inputs: xxx}

????????保存為 .ckpt 模型的一個(gè)好處是伟姐,當(dāng)需要繼續(xù)訓(xùn)練時(shí)收苏,只需要將訓(xùn)練過(guò)的模型結(jié)果導(dǎo)入,然后在這個(gè)基礎(chǔ)上再繼續(xù)訓(xùn)練愤兵。而下面的 .pb 格式則不能繼續(xù)訓(xùn)練鹿霸,因?yàn)檫@種格式保存的模型參數(shù)都已經(jīng)轉(zhuǎn)化為了常量(而不再是變量)。

二秆乳、pb 格式模型保存與恢復(fù)

????????.pb 格式模型保存與恢復(fù)相比于前面的 .ckpt 格式而言要稍微麻煩一點(diǎn)懦鼠,但使用更靈活,特別是模型恢復(fù),因?yàn)樗梢悦撾x會(huì)話(Session)而存在葛闷,便于部署憋槐。

1. pb 格式模型保存

????????與 .ckpt 格式模型保存類(lèi)似,首先定義數(shù)據(jù)入口淑趾、出口:

from tensorflow.python.framework import graph_util

···
inputs = tf.placeholder(tf.float32, shape=[None, ···], name='inputs') 
···
prediction = tf.nn.softmax(logits, name='prediction') 
···

with tf.Session() as sess:
    ···    <-- 訓(xùn)練過(guò)程
    graph_def = tf.get_default_graph().as_graph_def()
    output_graph_def = graph_util.convert_variables_to_constants(
        sess, 
        graph_def, 
        ['prediction']  <-- 參數(shù):output_node_names阳仔,輸出節(jié)點(diǎn)名
    )
    with tf.gfile.GFile('./xxx/xxx.pb', 'wb') as fid:
        serialized_graph = output_graph_def.SerializeToString()
        fid.write(serialized_graph)

然后通過(guò)函數(shù) graph_util.convert_variables_to_constants 將模型固話,使得所有變量轉(zhuǎn)化為常量扣泊,之后寫(xiě)入到指定的路徑完成模型保存過(guò)程近范。

2. pb 格式模型恢復(fù)

????????.pb 格式模型恢復(fù)自由度較大,不需要在會(huì)話里進(jìn)行操作延蟹,可以獨(dú)立存在:

import os

def load_model(path_to_model.pb):
    if not os.path.exists(path_to_model.pb):
        raise ValueError("'path_to_model.pb' is not exist.")

    model_graph = tf.Graph()
    with model_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(path_to_model.pb, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')
    return model_graph

模型導(dǎo)入之后评矩,便可以獲取數(shù)據(jù)入口和出口,然后進(jìn)行推斷:

model_graph = load_model('./xxx/xxx.pb')

inputs = model_graph.get_tensor_by_name('inputs:0')
prediction = model_graph.get_tensor_by_name('prediction:0')

with model_graph.as_default():
    with tf.Session(graph=model_graph) as sess:
        ···
        pred = sess.run(prediction, feed_dict={inputs: xxx}

三阱飘、ckpt 格式轉(zhuǎn) pb 格式

????????一般情況下斥杜,為了便于從斷點(diǎn)之處繼續(xù)訓(xùn)練,模型通常保存為 .ckpt 格式沥匈,而一旦對(duì)訓(xùn)練結(jié)果很滿(mǎn)意之后則可能需要將 .ckpt 格式轉(zhuǎn)化為 .pb 格式蔗喂。轉(zhuǎn)化方法很簡(jiǎn)單,只需要綜合前面的一高帖、二兩步即可:

from tensorflow.python.framework import graph_util

with tf.Session() as sess:
    # Load .ckpt file
    saver = tf.train.import_meta_graph('./xxx/xxx.ckpt.meta')
    saver.restore(sess, './xxx/xxx.ckpt')

    # Save as .pb file
    graph_def = tf.get_default_graph().as_graph_def()
    output_graph_def = graph_util.convert_variables_to_constants(
        sess, 
        graph_def, 
        ['prediction']  <-- 輸出節(jié)點(diǎn)名缰儿,以實(shí)際情況為準(zhǔn)
    )
    with tf.gfile.GFile('./xxx/xxx.pb', 'wb') as fid:
        serialized_graph = output_graph_def.SerializeToString()
        fid.write(serialized_graph)

????????預(yù)告:下一篇文章將簡(jiǎn)單介紹 tensorflow.contrib.slim 的應(yīng)用,敬請(qǐng)關(guān)注散址!

最后編輯于
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
  • 序言:七十年代末乖阵,一起剝皮案震驚了整個(gè)濱河市,隨后出現(xiàn)的幾起案子预麸,更是在濱河造成了極大的恐慌瞪浸,老刑警劉巖,帶你破解...
    沈念sama閱讀 211,042評(píng)論 6 490
  • 序言:濱河連續(xù)發(fā)生了三起死亡事件吏祸,死亡現(xiàn)場(chǎng)離奇詭異默终,居然都是意外死亡,警方通過(guò)查閱死者的電腦和手機(jī)犁罩,發(fā)現(xiàn)死者居然都...
    沈念sama閱讀 89,996評(píng)論 2 384
  • 文/潘曉璐 我一進(jìn)店門(mén),熙熙樓的掌柜王于貴愁眉苦臉地迎上來(lái)两疚,“玉大人床估,你說(shuō)我怎么就攤上這事∮詹常” “怎么了丐巫?”我有些...
    開(kāi)封第一講書(shū)人閱讀 156,674評(píng)論 0 345
  • 文/不壞的土叔 我叫張陵,是天一觀的道長(zhǎng)。 經(jīng)常有香客問(wèn)我递胧,道長(zhǎng)碑韵,這世上最難降的妖魔是什么? 我笑而不...
    開(kāi)封第一講書(shū)人閱讀 56,340評(píng)論 1 283
  • 正文 為了忘掉前任缎脾,我火速辦了婚禮祝闻,結(jié)果婚禮上,老公的妹妹穿的比我還像新娘遗菠。我一直安慰自己联喘,他們只是感情好,可當(dāng)我...
    茶點(diǎn)故事閱讀 65,404評(píng)論 5 384
  • 文/花漫 我一把揭開(kāi)白布辙纬。 她就那樣靜靜地躺著豁遭,像睡著了一般。 火紅的嫁衣襯著肌膚如雪贺拣。 梳的紋絲不亂的頭發(fā)上蓖谢,一...
    開(kāi)封第一講書(shū)人閱讀 49,749評(píng)論 1 289
  • 那天,我揣著相機(jī)與錄音譬涡,去河邊找鬼闪幽。 笑死,一個(gè)胖子當(dāng)著我的面吹牛昂儒,可吹牛的內(nèi)容都是我干的沟使。 我是一名探鬼主播,決...
    沈念sama閱讀 38,902評(píng)論 3 405
  • 文/蒼蘭香墨 我猛地睜開(kāi)眼渊跋,長(zhǎng)吁一口氣:“原來(lái)是場(chǎng)噩夢(mèng)啊……” “哼腊嗡!你這毒婦竟也來(lái)了?” 一聲冷哼從身側(cè)響起拾酝,我...
    開(kāi)封第一講書(shū)人閱讀 37,662評(píng)論 0 266
  • 序言:老撾萬(wàn)榮一對(duì)情侶失蹤燕少,失蹤者是張志新(化名)和其女友劉穎,沒(méi)想到半個(gè)月后蒿囤,有當(dāng)?shù)厝嗽跇?shù)林里發(fā)現(xiàn)了一具尸體客们,經(jīng)...
    沈念sama閱讀 44,110評(píng)論 1 303
  • 正文 獨(dú)居荒郊野嶺守林人離奇死亡,尸身上長(zhǎng)有42處帶血的膿包…… 初始之章·張勛 以下內(nèi)容為張勛視角 年9月15日...
    茶點(diǎn)故事閱讀 36,451評(píng)論 2 325
  • 正文 我和宋清朗相戀三年材诽,在試婚紗的時(shí)候發(fā)現(xiàn)自己被綠了底挫。 大學(xué)時(shí)的朋友給我發(fā)了我未婚夫和他白月光在一起吃飯的照片。...
    茶點(diǎn)故事閱讀 38,577評(píng)論 1 340
  • 序言:一個(gè)原本活蹦亂跳的男人離奇死亡脸侥,死狀恐怖建邓,靈堂內(nèi)的尸體忽然破棺而出,到底是詐尸還是另有隱情睁枕,我是刑警寧澤官边,帶...
    沈念sama閱讀 34,258評(píng)論 4 328
  • 正文 年R本政府宣布沸手,位于F島的核電站,受9級(jí)特大地震影響注簿,放射性物質(zhì)發(fā)生泄漏契吉。R本人自食惡果不足惜,卻給世界環(huán)境...
    茶點(diǎn)故事閱讀 39,848評(píng)論 3 312
  • 文/蒙蒙 一诡渴、第九天 我趴在偏房一處隱蔽的房頂上張望捐晶。 院中可真熱鬧,春花似錦玩徊、人聲如沸租悄。這莊子的主人今日做“春日...
    開(kāi)封第一講書(shū)人閱讀 30,726評(píng)論 0 21
  • 文/蒼蘭香墨 我抬頭看了看天上的太陽(yáng)泣棋。三九已至,卻和暖如春畔塔,著一層夾襖步出監(jiān)牢的瞬間潭辈,已是汗流浹背。 一陣腳步聲響...
    開(kāi)封第一講書(shū)人閱讀 31,952評(píng)論 1 264
  • 我被黑心中介騙來(lái)泰國(guó)打工澈吨, 沒(méi)想到剛下飛機(jī)就差點(diǎn)兒被人妖公主榨干…… 1. 我叫王不留把敢,地道東北人。 一個(gè)月前我還...
    沈念sama閱讀 46,271評(píng)論 2 360
  • 正文 我出身青樓谅辣,卻偏偏與公主長(zhǎng)得像修赞,于是被迫代替她去往敵國(guó)和親。 傳聞我的和親對(duì)象是個(gè)殘疾皇子桑阶,可洞房花燭夜當(dāng)晚...
    茶點(diǎn)故事閱讀 43,452評(píng)論 2 348

推薦閱讀更多精彩內(nèi)容