python 解析pb文件_tensorflow从ckpt和从.pb文件读取变量的值方式

最近在学习tensorflow自带的量化工具的相关知识，其中遇到的一个问题是从tensorflow保存好的ckpt文件或者是保存后的.pb文件(这里的pb是把权重和模型保存在一起的pb文件)读取权重，查看量化后的权重是否变成整形。

因此将自己解决这个问题记录下来，为了下一次遇到时，可以有所参考，也希望给有需要的同学一个可能的参考。

(1) 从保存的ckpt读取变量的值(以读取保存的第一个权重为例)

from tensorflow.python import pywrap_tensorflow

import tensorflow as tf

with tf.Graph().as_default():

with tf.Session() as sess:

ckpt = tf.train.get_checkpoint_state('./model_ckpt') #保存ckpt文件的文件夹

if ckpt and ckpt.model_checkpoint_path:

reader = pywrap_tensorflow.NewCheckpointReader('./model_ckpt/model.ckpt-999') #自己保存的ckpt文件名

all_variables = reader.get_variable_to_shape_map()

w1 = reader.get_tensor("Variable_1")

print(w1.shape)

print(w1)

else: print('No checkpoint file found')

(2) 从保存的.pb文件读取变量的值(以读取保存的第一个权重为例)

import tensorflow as tf

from tensorflow.python.framework import graph_util

from tensorflow.python.platform import gfile

import numpy as np

sess = tf.Session()

with gfile.FastGFile('Yourpb.pb', 'rb') as f: #自己保存的pb文件

graph_def = tf.GraphDef()

graph_def.ParseFromString(f.read())

sess.graph.as_default()

tf.import_graph_def(graph_def, name='')

print(sess.run('Variable_1:0'))

补充知识：如何从已存在的检查点文件(cpkt文件)种解析出里面变量——无需重新创建原始计算图

import tensorflow as tf

import os

CheckpointReader

tf.train.NewCheckpointReader是一个创建检查点读取器(CheckpointReader)对象的完美手段。 CheckpointReader中有几个非常有用的方法：

get_variable_to_shape_map() - 提供具有变量名称和形状的字典

debug_string() - 提供由检查点文件中所有变量组成的字符串

has_tensor(var_name) - 允许检查变量是否存在于检查点中

get_tensor(var_name) - 返回变量名称的张量

为了便于说明，我将定义一个函数来检查路径的有效性，并为您加载检查点读取器。

In [3]:

def load_reader(path):

assert os.path.exists(path), "Provided incorrect path to the file. {} doesn't exist".format(path)

return tf.train.NewCheckpointReader(path)

In [34]:

your_path = 'logs/squeezeDet1024x1024/train/model.ckpt-0'

reader = load_reader(your_path)

reader.debug_string()

用于返回包含以下内容的一个字符串：

variable name(变量名)

data type(数据类型)

tensor shape(张量类型)

它返回字符串的各元素间均用空格符' '分隔，你可以使用debug_string来创建一个变量名列表，如下所示：

In [53]:

all_var_descriptions = reader.debug_string().split()

var_names, var_shapes = all_var[::3], all_var[2::3]

print(var_names[:4])

print(var_shapes[:4])

输出：

['iou', 'fire9/squeeze1x1/kernels', 'fire9/squeeze1x1/biases', 'fire9/expand3x3/kernels/Momentum']

['[10,36864]', '[1,1,512,64]', '[64]', '[3,3,64,256]']

但是，对于完成同样的任务，更好的方法是使用reader.get_variable_to_shape_map()

reader.get_variable_to_shape_map()

用于返回包含所有变量及其形状名称的字典，变量作为字典的Key，形状作为Value。

In [66]:

saved_shapes = reader.get_variable_to_shape_map()

print('fire9/squeeze1x1/kernels:', saved_shapes['fire9/squeeze1x1/kernels'])

fire9/squeeze1x1/kernels: [1, 1, 512, 64]

reader.has_tensor(var_name)

返回bool值

这是一种方便的方法，允许您检查ckeckpoint中是否存在相关的变量。

In [51]:

names_that_exit = {var_name: reader.has_tensor(var_name) for var_name in var_names[:10]}

for key in names_that_exit:

print(key.decode()+':', names_that_exit[key])

fire8/squeeze1x1/kernels/Momentum: True

fire9/expand3x3/kernels: True

iou: True

fire9/expand3x3/biases: True

fire9/expand1x1/kernels: True

fire9/expand3x3/kernels/Momentum: True

fire9/expand1x1/biases/Momentum: True

fire9/squeeze1x1/biases: True

fire9/expand1x1/kernels/Momentum: True

fire9/squeeze1x1/kernels: True

reader.get_tensor(tensor_name)

返回包含检查点的张量值的NumPy数组

正常使用方法是先恢复一个张量，然后用恢复的张量初始化你自己的变量：

In [60]:

def recover_var(reader, var_name):

recovered_var = 'var to be recovered'

try:

recovered_var = reader.get_tensor(var_name)

except:

assert reader.has_tensor(var_name),\

"{} variable doesn't exist in the check point. Please check the variable name".format(var_name)

return recovered_var

In [67]:

checkpoint_var = recover_var(reader, 'conv1/kernels')

print ("Recovered variable has the following shape: \n", checkpoint_var.shape)

new_var = tf.Variable(initial_value=checkpoint_var, name="new_conv1")

print ("New variable will be initialized with recovered values and the following shape: \n", new_var.get_shape())

Recovered variable has the following shape:

(3, 3, 3, 64)

New variable will be initialized with recovered values and the following shape:

(3, 3, 3, 64)

以上这篇tensorflow从ckpt和从.pb文件读取变量的值方式就是小编分享给大家的全部内容了，希望能给大家一个参考，也希望大家多多支持脚本之家。