1. 程式人生 > 程式設計 >tensorflow從ckpt和從.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,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,64)
New variable will be initialized with recovered values and the following shape: 
(3,64)

以上這篇tensorflow從ckpt和從.pb檔案讀取變數的值方式就是小編分享給大家的全部內容了,希望能給大家一個參考,也希望大家多多支援我們。