ckpt轉pb

ckpt轉pb有兩種方式，一種是通過.meta文件加載圖和ckpt文件固化成ckpt，一種是加載運行一次網絡固化成pb，下面分別介紹

加載`meta`文件固化

加載.meta文件恢復網絡graph，然后加載ckpt，將變量轉換為constant常量，再移除ckpt中保留的訓練相關但和前向推理無關的結點。

import tensorflow as tf from tensorflow.python.framework import graph_utildef ckpt2pb():with tf.Graph().as_default() as graph_old:isess = tf.InteractiveSession()ckpt_filename = './model.ckpt'isess.run(tf.global_variables_initializer())saver = tf.train.import_meta_graph(ckpt_filename+'.meta', clear_devices=True)saver.restore(isess, ckpt_filename)constant_graph = graph_util.convert_variables_to_constants(isess, isess.graph_def, ["Cls/fc/biases"])constant_graph = graph_util.remove_training_nodes(constant_graph)with tf.gfile.GFile('./pb/model.pb', mode='wb') as f:f.write(constant_graph.SerializeToString())

運行一次網絡固化

運行網絡加載graph的優點是可以在網絡起止位置手動添加結點標記，方便定義網絡起止結點的名字。

def inference(is_training, img_input):'''your network'''def ckpt2pb2():with tf.Graph().as_default() as graph_old:img_input = tf.placeholder(tf.float32, shape=(None, 40, 120, 3))model = inference(False, img_input)isess = tf.InteractiveSession()ckpt_filename ='./model.ckpt'isess.run(tf.global_variables_initializer())saver = tf.train.Saver()saver.restore(isess, ckpt_filename)constant_graph = graph_util.convert_variables_to_constants(isess, isess.graph_def, ["Cls/fc/BiasAdd"])constant_graph = graph_util.remove_training_nodes(constant_graph)with tf.gfile.GFile('./pb/model.pb', mode='wb') as f:f.write(constant_graph.SerializeToString())

使用pb做前向推理

import tensorflow as tf import cv2 import numpy as npdef inference_use_pb():graph_path = './pb/module.pb'graph_def = tf.GraphDef()with tf.gfile.FastGFile(graph_path,'rb') as f:graph_def.ParseFromString(f.read())_ = importer.import_graph_def(graph_def, name="")isess = tf.InteractiveSession()images_placeholder = tf.get_default_graph().get_tensor_by_name("placeholder:0")embeddings = tf.get_default_graph().get_tensor_by_name("Cls/fc/BiasAdd:0")img = cv2.imread('./1.png')image = cv2.resize(img, (120, 40)).astype(np.float32) / 255.0image = np.reshape(image, (40,120,3))res = isess.run([embeddings], feed_dict={images_placeholder: np.reshape(image, [1, 40, 120, 3])})print(res)

讀取ckpt文件

有時我們需要從ckpt文件中獲取權重，則可以使用以下方法

import tensorflow as tfckpt_path = './model.ckpt'with tf.Session() as sess:for var_name, _ in tf.contrib.framework.list_variables(ckpt_path):print(var_name)var = tf.contrib.framework.load_variable(ckpt_path, var_name)print(var.shape)print(var)

讀取pb文件

當沒有ckpt文件，只有pb文件時，使用以下方法獲取權重

from tensorflow.python.platform import gfile from tensorflow.python.framework import tensor_utilgraph_path = './model.pb'def values_from_const(node_def):if node_def.op != "Const":raise ValueError("Node named '%s' should be a Const op for values_from_const." % node_def.name)input_tensor = node_def.attr["value"].tensortensor_value = tensor_util.MakeNdarray(input_tensor)return tensor_valuedef read_pb():input_graph_def = graph_pb2.GraphDef()with gfile.Open(graph_path, "rb") as f:data = f.read()input_graph_def.ParseFromString(data)for node in input_graph_def.node:print(node.name)print(node.op)if node.op == "Const":if 'weights' in node.name:weight = values_from_const(node).reshape(-1)print(weight)if __name__ == "__main__":read_pb()

總結

以上是生活随笔為你收集整理的tensor转list_tensorflow中ckpt转pb的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。