ckpt轉為pb

本博客實現將自己訓練保存的ckpt模型轉換為pb文件，該方法適用於任何ckpt模型，當然你需要確定ckpt模型輸入/輸出的節點名稱。

目錄

tensorflow實現將ckpt轉pb文件

一、CKPT 轉換成 PB格式

二、 pb模型預測

三、源碼下載和資料推薦

1、訓練方法

2、本博客Github地址

3、將模型移植Android的方法

使用 tf.train.saver()保存模型時會產生多個文件，會把計算圖的結構和圖上參數取值分成了不同的文件存儲。這種方法是在TensorFlow中是最常用的保存方式。

例如：下面的代碼運行后，會在save目錄下保存了四個文件：

import tensorflow as tf
# 聲明兩個變量
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的值一會讀取之后對比
print("v2:", sess.run(v2))
saver_path = saver.save(sess, "save/model.ckpt") # 將模型保存到save/model.ckpt文件
print("Model saved in file:", saver_path)

其中

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

TensoFlow為我們提供了convert_variables_to_constants()方法，該方法可以固化模型結構，將計算圖中的變量取值以常量的形式保存，而且保存的模型可以移植到Android平台。

一、CKPT 轉換成 PB格式
將CKPT 轉換成 PB格式的文件的過程可簡述如下：

通過傳入 CKPT 模型的路徑得到模型的圖和變量數據
通過 import_meta_graph 導入模型中的圖
通過 saver.restore 從模型中恢復圖中各個變量的數據
通過 graph_util.convert_variables_to_constants 將模型持久化
下面的CKPT 轉換成 PB格式例子，是我訓練GoogleNet InceptionV3模型保存的ckpt轉pb文件的例子，訓練過程可參考博客：

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

def freeze_graph(input_checkpoint,output_graph):
'''
:param input_checkpoint:
:param output_graph: PB模型保存路徑
:return:
'''
# checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態是否可用
# input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑

# 指定輸出的節點名稱,該節點名稱必須是原模型中存在的節點
output_node_names = "InceptionV3/Logits/SpatialSqueeze"
saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
graph = tf.get_default_graph() # 獲得默認的圖
input_graph_def = graph.as_graph_def() # 返回一個序列化的圖代表當前的圖

with tf.Session() as sess:
saver.restore(sess, input_checkpoint) #恢復圖並得到數據
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(","))# 如果有多個輸出節點，以逗號隔開

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)) #得到當前圖有幾個操作節點

# for op in graph.get_operations():
# print(op.name, op.values())
說明：

1、函數freeze_graph中，最重要的就是要確定“指定輸出的節點名稱”，這個節點名稱必須是原模型中存在的節點，對於freeze操作，我們需要定義輸出結點的名字。因為網絡其實是比較復雜的，定義了輸出結點的名字，那么freeze的時候就只把輸出該結點所需要的子圖都固化下來，其他無關的就舍棄掉。因為我們freeze模型的目的是接下來做預測。所以，output_node_names一般是網絡模型最后一層輸出的節點名稱，或者說就是我們預測的目標。

2、在保存的時候，通過convert_variables_to_constants函數來指定需要固化的節點名稱，對於鄙人的代碼，需要固化的節點只有一個：output_node_names。注意節點名稱與張量的名稱的區別，例如：“input:0”是張量的名稱，而"input"表示的是節點的名稱。

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

4、實質上，我們可以直接在恢復的會話sess中，獲得默認的網絡圖，更簡單的方法，如下：

def freeze_graph(input_checkpoint,output_graph):
'''
:param input_checkpoint:
:param output_graph: PB模型保存路徑
:return:
'''
# checkpoint = tf.train.get_checkpoint_state(model_folder) #檢查目錄下ckpt文件狀態是否可用
# input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑

# 指定輸出的節點名稱,該節點名稱必須是原模型中存在的節點
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) #恢復圖並得到數據
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(","))# 如果有多個輸出節點，以逗號隔開

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)) #得到當前圖有幾個操作節點
調用方法很簡單，輸入ckpt模型路徑，輸出pb模型的路徑即可：

# 輸入ckpt模型路徑
input_checkpoint='models/model.ckpt-10000'
# 輸出pb模型的路徑
out_pb_path="models/pb/frozen_model.pb"
# 調用freeze_graph將ckpt轉為pb
freeze_graph(input_checkpoint,out_pb_path)
5、上面以及說明：在保存的時候，通過convert_variables_to_constants函數來指定需要固化的節點名稱，對於鄙人的代碼，需要固化的節點只有一個：output_node_names。因此，其他網絡模型，也可以通過簡單的修改輸出的節點名稱output_node_names，將ckpt轉為pb文件 。

PS：注意節點名稱，應包含name_scope 和 variable_scope命名空間，並用“/”隔開，如"InceptionV3/Logits/SpatialSqueeze"

二、 pb模型預測
下面是預測pb模型的代碼

def freeze_graph_test(pb_path, image_path):
'''
:param pb_path:pb文件的路徑
:param image_path:測試圖片的路徑
: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())

# 定義輸入的張量名稱,對應網絡結構的輸入張量
# input:0作為輸入圖像,keep_prob:0作為dropout的參數,測試時值為1,is_training:0訓練參數
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")

# 讀取測試圖片
im=read_image(image_path,resize_height,resize_width,normalization=True)
im=im[np.newaxis,:]
# 測試讀出來的模型是否正確，注意這里傳入的是輸出和輸入節點的tensor的名字，不是操作節點的名字
# 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預測不同的是，pb文件已經固化了網絡模型結構，因此，即使不知道原訓練模型（train）的源碼，我們也可以恢復網絡圖，並進行預測。恢復模型十分簡單，只需要從讀取的序列化數據中導入網絡結構即可：

tf.import_graph_def(output_graph_def, name="")
2、但必須知道原網絡模型的輸入和輸出的節點名稱（當然了，傳遞數據時，是通過輸入輸出的張量來完成的）。由於InceptionV3模型的輸入有三個節點，因此這里需要定義輸入的張量名稱，它對應網絡結構的輸入張量：

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、預測時，需要feed輸入數據：

# 測試讀出來的模型是否正確，注意這里傳入的是輸出和輸入節點的tensor的名字，不是操作節點的名字
# 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、其他網絡模型預測時，也可以通過修改輸入和輸出的張量的名稱 。

PS：注意張量的名稱，即為：節點名稱+“:”+“id號”，如"InceptionV3/Logits/SpatialSqueeze:0"

完整的CKPT 轉換成 PB格式和預測的代碼如下：

# -*-coding: utf-8 -*-
"""
@Project: tensorflow_models_nets
@File : convert_pb.py
@Author : panjq
@E-mail : [e-mail]
@Date : 2018-08-29 17:46:50
@info :
-通過傳入 CKPT 模型的路徑得到模型的圖和變量數據
-通過 import_meta_graph 導入模型中的圖
-通過 saver.restore 從模型中恢復圖中各個變量的數據
-通過 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:測試圖片的路徑
: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())

# 定義輸入的張量名稱,對應網絡結構的輸入張量
# input:0作為輸入圖像,keep_prob:0作為dropout的參數,測試時值為1,is_training:0訓練參數
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")

# 讀取測試圖片
im=read_image(image_path,resize_height,resize_width,normalization=True)
im=im[np.newaxis,:]
# 測試讀出來的模型是否正確，注意這里傳入的是輸出和輸入節點的tensor的名字，不是操作節點的名字
# 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文件狀態是否可用
# input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑

# 指定輸出的節點名稱,該節點名稱必須是原模型中存在的節點
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) #恢復圖並得到數據
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(","))# 如果有多個輸出節點，以逗號隔開

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)) #得到當前圖有幾個操作節點

# 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文件狀態是否可用
# input_checkpoint = checkpoint.model_checkpoint_path #得ckpt文件路徑

# 指定輸出的節點名稱,該節點名稱必須是原模型中存在的節點
output_node_names = "InceptionV3/Logits/SpatialSqueeze"
saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
graph = tf.get_default_graph() # 獲得默認的圖
input_graph_def = graph.as_graph_def() # 返回一個序列化的圖代表當前的圖

with tf.Session() as sess:
saver.restore(sess, input_checkpoint) #恢復圖並得到數據
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(","))# 如果有多個輸出節點，以逗號隔開

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)) #得到當前圖有幾個操作節點

# 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"
# 調用freeze_graph將ckpt轉為pb
freeze_graph(input_checkpoint,out_pb_path)

# 測試pb模型
image_path = 'test_image/animal.jpg'
freeze_graph_test(pb_path=out_pb_path, image_path=image_path)
三、源碼下載和資料推薦
1、訓練方法
上面的CKPT 轉換成 PB格式例子，是我訓練GoogleNet InceptionV3模型保存的ckpt轉pb文件的例子，訓練過程可參考博客：

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

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

預訓練模型下載地址：https://download.csdn.net/download/guyuealian/10610847

3、將模型移植Android的方法
pb文件是可以移植到Android平台運行的，其方法，可參考：

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

https://blog.csdn.net/guyuealian/article/details/79672257
————————————————

原文鏈接：https://blog.csdn.net/guyuealian/article/details/82218092