快捷導(dǎo)航

tensorflow實(shí)現(xiàn)將ckpt轉(zhuǎn)pb文件的方法

更新時(shí)間：2020年04月22日 11:14:10 作者：努力努力再努力tq

這篇文章主要介紹了tensorflow實(shí)現(xiàn)將ckpt轉(zhuǎn)pb文件的方法，本文通過實(shí)例代碼給大家介紹的非常詳細(xì)，對(duì)大家的學(xué)習(xí)或工作具有一定的參考借鑒價(jià)值,需要的朋友可以參考下

本博客實(shí)現(xiàn)將自己訓(xùn)練保存的ckpt模型轉(zhuǎn)換為pb文件，該方法適用于任何ckpt模型，當(dāng)然你需要確定ckpt模型輸入/輸出的節(jié)點(diǎn)名稱。

使用 tf.train.saver()保存模型時(shí)會(huì)產(chǎn)生多個(gè)文件，會(huì)把計(jì)算圖的結(jié)構(gòu)和圖上參數(shù)取值分成了不同的文件存儲(chǔ)。這種方法是在TensorFlow中是最常用的保存方式。

例如：下面的代碼運(yùn)行后，會(huì)在save目錄下保存了四個(gè)文件：

import tensorflow as tf
# 聲明兩個(gè)變量
v1 = tf.Variable(tf.random_normal([1, 2]), name="v1")
v2 = tf.Variable(tf.random_normal([2, 3]), name="v2")
init_op = tf.global_variables_initializer() # 初始化全部變量
saver = tf.train.Saver() # 聲明tf.train.Saver類用于保存模型
with tf.Session() as sess:
 sess.run(init_op)
 print("v1:", sess.run(v1)) # 打印v1、v2的值一會(huì)讀取之后對(duì)比
 print("v2:", sess.run(v2))
 saver_path = saver.save(sess, "save/model.ckpt") # 將模型保存到save/model.ckpt文件
 print("Model saved in file:", saver_path)

其中,checkpoint是檢查點(diǎn)文件，文件保存了一個(gè)目錄下所有的模型文件列表；
model.ckpt.meta文件保存了TensorFlow計(jì)算圖的結(jié)構(gòu)，可以理解為神經(jīng)網(wǎng)絡(luò)的網(wǎng)絡(luò)結(jié)構(gòu)，該文件可以被 tf.train.import_meta_graph 加載到當(dāng)前默認(rèn)的圖來使用。
ckpt.data ：保存模型中每個(gè)變量的取值
但很多時(shí)候，我們需要將TensorFlow的模型導(dǎo)出為單個(gè)文件（同時(shí)包含模型結(jié)構(gòu)的定義與權(quán)重），方便在其他地方使用（如在Android中部署網(wǎng)絡(luò)）。利用tf.train.write_graph()默認(rèn)情況下只導(dǎo)出了網(wǎng)絡(luò)的定義（沒有權(quán)重），而利用tf.train.Saver().save()導(dǎo)出的文件graph_def與權(quán)重是分離的，因此需要采用別的方法。我們知道，graph_def文件中沒有包含網(wǎng)絡(luò)中的Variable值（通常情況存儲(chǔ)了權(quán)重），但是卻包含了constant值，所以如果我們能把Variable轉(zhuǎn)換為constant，即可達(dá)到使用一個(gè)文件同時(shí)存儲(chǔ)網(wǎng)絡(luò)架構(gòu)與權(quán)重的目標(biāo)。

TensoFlow為我們提供了convert_variables_to_constants()方法，該方法可以固化模型結(jié)構(gòu)，將計(jì)算圖中的變量取值以常量的形式保存，而且保存的模型可以移植到Android平臺(tái)。

一、CKPT 轉(zhuǎn)換成 PB格式

將CKPT 轉(zhuǎn)換成 PB格式的文件的過程可簡(jiǎn)述如下：

通過傳入 CKPT 模型的路徑得到模型的圖和變量數(shù)據(jù)
通過 import_meta_graph 導(dǎo)入模型中的圖
通過 saver.restore 從模型中恢復(fù)圖中各個(gè)變量的數(shù)據(jù)
通過 graph_util.convert_variables_to_constants 將模型持久化
下面的CKPT 轉(zhuǎn)換成 PB格式例子，是我訓(xùn)練GoogleNet InceptionV3模型保存的ckpt轉(zhuǎn)pb文件的例子，訓(xùn)練過程可參考博客：《使用自己的數(shù)據(jù)集訓(xùn)練GoogLenet InceptionNet V1 V2 V3模型（TensorFlow）》：

def freeze_graph(input_checkpoint,output_graph):
 '''
 :param input_checkpoint:
 :param output_graph: PB模型保存路徑
 :return:
 '''
 # checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態(tài)是否可用
 # input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
 
 # 指定輸出的節(jié)點(diǎn)名稱,該節(jié)點(diǎn)名稱必須是原模型中存在的節(jié)點(diǎn)
 output_node_names = "InceptionV3/Logits/SpatialSqueeze"
 saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
 graph = tf.get_default_graph() # 獲得默認(rèn)的圖
 input_graph_def = graph.as_graph_def() # 返回一個(gè)序列化的圖代表當(dāng)前的圖
 
 with tf.Session() as sess:
 saver.restore(sess, input_checkpoint) #恢復(fù)圖并得到數(shù)據(jù)
 output_graph_def = graph_util.convert_variables_to_constants( # 模型持久化，將變量值固定
 sess=sess,
 input_graph_def=input_graph_def,# 等于:sess.graph_def
 output_node_names=output_node_names.split(","))# 如果有多個(gè)輸出節(jié)點(diǎn)，以逗號(hào)隔開
 
 with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
 f.write(output_graph_def.SerializeToString()) #序列化輸出
 print("%d ops in the final graph." % len(output_graph_def.node)) #得到當(dāng)前圖有幾個(gè)操作節(jié)點(diǎn)
 
 # for op in graph.get_operations():
 # print(op.name, op.values())

說明：

1、函數(shù)freeze_graph中，最重要的就是要確定“指定輸出的節(jié)點(diǎn)名稱”，這個(gè)節(jié)點(diǎn)名稱必須是原模型中存在的節(jié)點(diǎn)，對(duì)于freeze操作，我們需要定義輸出結(jié)點(diǎn)的名字。因?yàn)榫W(wǎng)絡(luò)其實(shí)是比較復(fù)雜的，定義了輸出結(jié)點(diǎn)的名字，那么freeze的時(shí)候就只把輸出該結(jié)點(diǎn)所需要的子圖都固化下來，其他無關(guān)的就舍棄掉。因?yàn)槲覀僨reeze模型的目的是接下來做預(yù)測(cè)。所以，output_node_names一般是網(wǎng)絡(luò)模型最后一層輸出的節(jié)點(diǎn)名稱，或者說就是我們預(yù)測(cè)的目標(biāo)。

2、在保存的時(shí)候，通過convert_variables_to_constants函數(shù)來指定需要固化的節(jié)點(diǎn)名稱，對(duì)于鄙人的代碼，需要固化的節(jié)點(diǎn)只有一個(gè)：output_node_names。注意節(jié)點(diǎn)名稱與張量的名稱的區(qū)別，例如：“input:0”是張量的名稱，而"input"表示的是節(jié)點(diǎn)的名稱。

3、源碼中通過graph = tf.get_default_graph()獲得默認(rèn)的圖，這個(gè)圖就是由saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)恢復(fù)的圖，因此必須先執(zhí)行tf.train.import_meta_graph，再執(zhí)行tf.get_default_graph() 。

4、實(shí)質(zhì)上，我們可以直接在恢復(fù)的會(huì)話sess中，獲得默認(rèn)的網(wǎng)絡(luò)圖，更簡(jiǎn)單的方法，如下：

def freeze_graph(input_checkpoint,output_graph):
 '''
 :param input_checkpoint:
 :param output_graph: PB模型保存路徑
 :return:
 '''
 # checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態(tài)是否可用
 # input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
 
 # 指定輸出的節(jié)點(diǎn)名稱,該節(jié)點(diǎn)名稱必須是原模型中存在的節(jié)點(diǎn)
 output_node_names = "InceptionV3/Logits/SpatialSqueeze"
 saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
 
 with tf.Session() as sess:
 saver.restore(sess, input_checkpoint) #恢復(fù)圖并得到數(shù)據(jù)
 output_graph_def = graph_util.convert_variables_to_constants( # 模型持久化，將變量值固定
 sess=sess,
 input_graph_def=sess.graph_def,# 等于:sess.graph_def
 output_node_names=output_node_names.split(","))# 如果有多個(gè)輸出節(jié)點(diǎn)，以逗號(hào)隔開
 
 with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
 f.write(output_graph_def.SerializeToString()) #序列化輸出
 print("%d ops in the final graph." % len(output_graph_def.node)) #得到當(dāng)前圖有幾個(gè)操作節(jié)點(diǎn)

調(diào)用方法很簡(jiǎn)單，輸入ckpt模型路徑，輸出pb模型的路徑即可：

# 輸入ckpt模型路徑
input_checkpoint='models/model.ckpt-10000'
# 輸出pb模型的路徑
out_pb_path="models/pb/frozen_model.pb"
# 調(diào)用freeze_graph將ckpt轉(zhuǎn)為pb
freeze_graph(input_checkpoint,out_pb_path)

5、上面以及說明：在保存的時(shí)候，通過convert_variables_to_constants函數(shù)來指定需要固化的節(jié)點(diǎn)名稱，對(duì)于鄙人的代碼，需要固化的節(jié)點(diǎn)只有一個(gè)：output_node_names。因此，其他網(wǎng)絡(luò)模型，也可以通過簡(jiǎn)單的修改輸出的節(jié)點(diǎn)名稱output_node_names，將ckpt轉(zhuǎn)為pb文件。

PS：注意節(jié)點(diǎn)名稱，應(yīng)包含name_scope 和 variable_scope命名空間，并用“/”隔開，如"InceptionV3/Logits/SpatialSqueeze"

二、 pb模型預(yù)測(cè)

下面是預(yù)測(cè)pb模型的代碼

def freeze_graph_test(pb_path, image_path):
 '''
 :param pb_path:pb文件的路徑
 :param image_path:測(cè)試圖片的路徑
 :return:
 '''
 with tf.Graph().as_default():
 output_graph_def = tf.GraphDef()
 with open(pb_path, "rb") as f:
 output_graph_def.ParseFromString(f.read())
 tf.import_graph_def(output_graph_def, name="")
 with tf.Session() as sess:
 sess.run(tf.global_variables_initializer())
 
 # 定義輸入的張量名稱,對(duì)應(yīng)網(wǎng)絡(luò)結(jié)構(gòu)的輸入張量
 # input:0作為輸入圖像,keep_prob:0作為dropout的參數(shù),測(cè)試時(shí)值為1,is_training:0訓(xùn)練參數(shù)
 input_image_tensor = sess.graph.get_tensor_by_name("input:0")
 input_keep_prob_tensor = sess.graph.get_tensor_by_name("keep_prob:0")
 input_is_training_tensor = sess.graph.get_tensor_by_name("is_training:0")
 
 # 定義輸出的張量名稱
 output_tensor_name = sess.graph.get_tensor_by_name("InceptionV3/Logits/SpatialSqueeze:0")
 
 # 讀取測(cè)試圖片
 im=read_image(image_path,resize_height,resize_width,normalization=True)
 im=im[np.newaxis,:]
 # 測(cè)試讀出來的模型是否正確，注意這里傳入的是輸出和輸入節(jié)點(diǎn)的tensor的名字，不是操作節(jié)點(diǎn)的名字
 # out=sess.run("InceptionV3/Logits/SpatialSqueeze:0", feed_dict={'input:0': im,'keep_prob:0':1.0,'is_training:0':False})
 out=sess.run(output_tensor_name, feed_dict={input_image_tensor: im,
 input_keep_prob_tensor:1.0,
 input_is_training_tensor:False})
 print("out:{}".format(out))
 score = tf.nn.softmax(out, name='pre')
 class_id = tf.argmax(score, 1)
 print "pre class_id:{}".format(sess.run(class_id))

說明：

1、與ckpt預(yù)測(cè)不同的是，pb文件已經(jīng)固化了網(wǎng)絡(luò)模型結(jié)構(gòu)，因此，即使不知道原訓(xùn)練模型（train）的源碼，我們也可以恢復(fù)網(wǎng)絡(luò)圖，并進(jìn)行預(yù)測(cè)。恢復(fù)模型十分簡(jiǎn)單，只需要從讀取的序列化數(shù)據(jù)中導(dǎo)入網(wǎng)絡(luò)結(jié)構(gòu)即可：

tf.import_graph_def(output_graph_def, name="")
2、但必須知道原網(wǎng)絡(luò)模型的輸入和輸出的節(jié)點(diǎn)名稱（當(dāng)然了，傳遞數(shù)據(jù)時(shí)，是通過輸入輸出的張量來完成的）。由于InceptionV3模型的輸入有三個(gè)節(jié)點(diǎn)，因此這里需要定義輸入的張量名稱，它對(duì)應(yīng)網(wǎng)絡(luò)結(jié)構(gòu)的輸入張量：

input_image_tensor = sess.graph.get_tensor_by_name("input:0")
input_keep_prob_tensor = sess.graph.get_tensor_by_name("keep_prob:0")
input_is_training_tensor = sess.graph.get_tensor_by_name("is_training:0")
以及輸出的張量名稱：

output_tensor_name = sess.graph.get_tensor_by_name("InceptionV3/Logits/SpatialSqueeze:0")

3、預(yù)測(cè)時(shí)，需要feed輸入數(shù)據(jù)：

# 測(cè)試讀出來的模型是否正確，注意這里傳入的是輸出和輸入節(jié)點(diǎn)的tensor的名字，不是操作節(jié)點(diǎn)的名字
# out=sess.run("InceptionV3/Logits/SpatialSqueeze:0", feed_dict={'input:0': im,'keep_prob:0':1.0,'is_training:0':False})
out=sess.run(output_tensor_name, feed_dict={input_image_tensor: im,
input_keep_prob_tensor:1.0,
input_is_training_tensor:False})

4、其他網(wǎng)絡(luò)模型預(yù)測(cè)時(shí)，也可以通過修改輸入和輸出的張量的名稱。

PS：注意張量的名稱，即為：節(jié)點(diǎn)名稱+“:”+“id號(hào)”，如"InceptionV3/Logits/SpatialSqueeze:0"

完整的CKPT 轉(zhuǎn)換成 PB格式和預(yù)測(cè)的代碼如下：

# -*-coding: utf-8 -*-
"""
 @Project: tensorflow_models_nets
 @File : convert_pb.py
 @Author : panjq
 @E-mail : pan_jinquan@163.com
 @Date : 2018-08-29 17:46:50
 @info :
 -通過傳入 CKPT 模型的路徑得到模型的圖和變量數(shù)據(jù)
 -通過 import_meta_graph 導(dǎo)入模型中的圖
 -通過 saver.restore 從模型中恢復(fù)圖中各個(gè)變量的數(shù)據(jù)
 -通過 graph_util.convert_variables_to_constants 將模型持久化
"""
 
import tensorflow as tf
from create_tf_record import *
from tensorflow.python.framework import graph_util
 
resize_height = 299 # 指定圖片高度
resize_width = 299 # 指定圖片寬度
depths = 3
 
def freeze_graph_test(pb_path, image_path):
 '''
 :param pb_path:pb文件的路徑
 :param image_path:測(cè)試圖片的路徑
 :return:
 '''
 with tf.Graph().as_default():
 output_graph_def = tf.GraphDef()
 with open(pb_path, "rb") as f:
 output_graph_def.ParseFromString(f.read())
 tf.import_graph_def(output_graph_def, name="")
 with tf.Session() as sess:
 sess.run(tf.global_variables_initializer())
 
 # 定義輸入的張量名稱,對(duì)應(yīng)網(wǎng)絡(luò)結(jié)構(gòu)的輸入張量
 # input:0作為輸入圖像,keep_prob:0作為dropout的參數(shù),測(cè)試時(shí)值為1,is_training:0訓(xùn)練參數(shù)
 input_image_tensor = sess.graph.get_tensor_by_name("input:0")
 input_keep_prob_tensor = sess.graph.get_tensor_by_name("keep_prob:0")
 input_is_training_tensor = sess.graph.get_tensor_by_name("is_training:0")
 
 # 定義輸出的張量名稱
 output_tensor_name = sess.graph.get_tensor_by_name("InceptionV3/Logits/SpatialSqueeze:0")
 
 # 讀取測(cè)試圖片
 im=read_image(image_path,resize_height,resize_width,normalization=True)
 im=im[np.newaxis,:]
 # 測(cè)試讀出來的模型是否正確，注意這里傳入的是輸出和輸入節(jié)點(diǎn)的tensor的名字，不是操作節(jié)點(diǎn)的名字
 # out=sess.run("InceptionV3/Logits/SpatialSqueeze:0", feed_dict={'input:0': im,'keep_prob:0':1.0,'is_training:0':False})
 out=sess.run(output_tensor_name, feed_dict={input_image_tensor: im,
 input_keep_prob_tensor:1.0,
 input_is_training_tensor:False})
 print("out:{}".format(out))
 score = tf.nn.softmax(out, name='pre')
 class_id = tf.argmax(score, 1)
 print "pre class_id:{}".format(sess.run(class_id))
 
 
def freeze_graph(input_checkpoint,output_graph):
 '''
 :param input_checkpoint:
 :param output_graph: PB模型保存路徑
 :return:
 '''
 # checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態(tài)是否可用
 # input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
 
 # 指定輸出的節(jié)點(diǎn)名稱,該節(jié)點(diǎn)名稱必須是原模型中存在的節(jié)點(diǎn)
 output_node_names = "InceptionV3/Logits/SpatialSqueeze"
 saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
 
 with tf.Session() as sess:
 saver.restore(sess, input_checkpoint) #恢復(fù)圖并得到數(shù)據(jù)
 output_graph_def = graph_util.convert_variables_to_constants( # 模型持久化，將變量值固定
 sess=sess,
 input_graph_def=sess.graph_def,# 等于:sess.graph_def
 output_node_names=output_node_names.split(","))# 如果有多個(gè)輸出節(jié)點(diǎn)，以逗號(hào)隔開
 
 with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
 f.write(output_graph_def.SerializeToString()) #序列化輸出
 print("%d ops in the final graph." % len(output_graph_def.node)) #得到當(dāng)前圖有幾個(gè)操作節(jié)點(diǎn)
 
 # for op in sess.graph.get_operations():
 # print(op.name, op.values())
 
def freeze_graph2(input_checkpoint,output_graph):
 '''
 :param input_checkpoint:
 :param output_graph: PB模型保存路徑
 :return:
 '''
 # checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態(tài)是否可用
 # input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑
 
 # 指定輸出的節(jié)點(diǎn)名稱,該節(jié)點(diǎn)名稱必須是原模型中存在的節(jié)點(diǎn)
 output_node_names = "InceptionV3/Logits/SpatialSqueeze"
 saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
 graph = tf.get_default_graph() # 獲得默認(rèn)的圖
 input_graph_def = graph.as_graph_def() # 返回一個(gè)序列化的圖代表當(dāng)前的圖
 
 with tf.Session() as sess:
 saver.restore(sess, input_checkpoint) #恢復(fù)圖并得到數(shù)據(jù)
 output_graph_def = graph_util.convert_variables_to_constants( # 模型持久化，將變量值固定
 sess=sess,
 input_graph_def=input_graph_def,# 等于:sess.graph_def
 output_node_names=output_node_names.split(","))# 如果有多個(gè)輸出節(jié)點(diǎn)，以逗號(hào)隔開
 
 with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
 f.write(output_graph_def.SerializeToString()) #序列化輸出
 print("%d ops in the final graph." % len(output_graph_def.node)) #得到當(dāng)前圖有幾個(gè)操作節(jié)點(diǎn)
 
 # for op in graph.get_operations():
 # print(op.name, op.values())
 
 
if __name__ == '__main__':
 # 輸入ckpt模型路徑
 input_checkpoint='models/model.ckpt-10000'
 # 輸出pb模型的路徑
 out_pb_path="models/pb/frozen_model.pb"
 # 調(diào)用freeze_graph將ckpt轉(zhuǎn)為pb
 freeze_graph(input_checkpoint,out_pb_path)
 
 # 測(cè)試pb模型
 image_path = 'test_image/animal.jpg'
 freeze_graph_test(pb_path=out_pb_path, image_path=image_path)

三、源碼下載和資料推薦

1、訓(xùn)練方法
上面的CKPT 轉(zhuǎn)換成 PB格式例子，是我訓(xùn)練GoogleNet InceptionV3模型保存的ckpt轉(zhuǎn)pb文件的例子，訓(xùn)練過程可參考博客：

《使用自己的數(shù)據(jù)集訓(xùn)練GoogLenet InceptionNet V1 V2 V3模型（TensorFlow）》：https://blog.csdn.net/guyuealian/article/details/81560537

2、Github地址
Github源碼：https://github.com/PanJinquan/tensorflow_models_nets 中的convert_pb.py文件

預(yù)訓(xùn)練模型下載地址：http://xiazai.jb51.net/202004/yuanma/googlenet_inception_jb51.rar

3、將模型移植Android的方法
pb文件是可以移植到Android平臺(tái)運(yùn)行的，其方法，可參考：

《將tensorflow訓(xùn)練好的模型移植到Android (MNIST手寫數(shù)字識(shí)別)》

參考：

[1] http://www.dbjr.com.cn/article/185209.htm

【2】http://www.dbjr.com.cn/article/185206.htm

到此這篇關(guān)于tensorflow實(shí)現(xiàn)將ckpt轉(zhuǎn)pb文件的方法的文章就介紹到這了,更多相關(guān)tensorflow ckpt轉(zhuǎn)pb文件內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: