tf34：从ckpt中读取权重值

MachineLP

发布于 2022-05-09 14:58:11

1.1K0

发布于 2022-05-09 14:58:11

文章被收录于专栏：小鹏的专栏

在TensorFlow里，提供了tf.train.NewCheckpointReader来查看model.ckpt文件中保存的变量信息。

一个简单的例子：

import tensorflow as tf
  
w = tf.Variable(2, dtype=tf.float32, name='w')  
b = tf.Variable(1, dtype=tf.float32, name='b')  

x = tf.placeholder(tf.float32, shape=[1], name='x')  
  
logit = w * x + b
  
init = tf.initialize_all_variables()  
  
saver = tf.train.Saver()  
  
with tf.Session() as sess:  
    sess.run(init)  
    saver.save(sess, "./model.ckpt")

import tensorflow as tf
  
reader = tf.train.NewCheckpointReader("./model.ckpt")  
  
variables = reader.get_variable_to_shape_map()  
  
for v in variables: 
    w = reader.get_tensor(v)  
    print(type(w))  
    # print(w.shape) 
    # print (w[0]) 
    print(w)