输出TensorFlow保存的checkpoint内变量的几种方法

最新推荐文章于 2023-09-26 19:01:30 发布

云游客来

最新推荐文章于 2023-09-26 19:01:30 发布

阅读量1.7k

点赞数

分类专栏： tensorflow

本文链接：https://blog.csdn.net/womnljf/article/details/81630869

版权

tensorflow 专栏收录该内容

5 篇文章 0 订阅

订阅专栏

1. tensorflow官方save and restore部分。

其中 inspect_checkpoint library,可以用来检查checkpoint中的变量。（https://www.tensorflow.org/guide/saved_model?hl=zh-cn）。

# import the inspect_checkpoint library
from tensorflow.python.tools import inspect_checkpoint as chkp

# print all tensors in checkpoint file
chkp.print_tensors_in_checkpoint_file("/tmp/model.ckpt", tensor_name='', all_tensors=True)

# tensor_name:  v1
# [ 1.  1.  1.]
# tensor_name:  v2
# [-1. -1. -1. -1. -1.]

# print only tensor v1 in checkpoint file
chkp.print_tensors_in_checkpoint_file("/tmp/model.ckpt", tensor_name='v1', all_tensors=False)

# tensor_name:  v1
# [ 1.  1.  1.]

# print only tensor v2 in checkpoint file
chkp.print_tensors_in_checkpoint_file("/tmp/model.ckpt", tensor_name='v2', all_tensors=False)

# tensor_name:  v2
# [-1. -1. -1. -1. -1.]

2. 第二种方法是使用 tensorflow/python/tools/inspect_checkpoint.py 中提到的tf.train.NewCheckpointReader类

from tensorflow.python import pywrap_tensorflow
checkpoint_path = os.path.join(model_dir, "model.ckpt")
reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
for key in var_to_shape_map:
    print("tensor_name: ", key)
    print(reader.get_tensor(key)) # Remove this is you want to print only variable names