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'))