tensor转list_tensorflow中ckpt转pb

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的全部内容，希望文章能够帮你解决所遇到的问题。

如果觉得生活随笔网站内容还不错，欢迎将生活随笔推荐给好友。