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'))
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章