Tensorflow 读取ckpt文件中的tensor操作

时间:2022-03-08 19:18:53

在使用pre-train model时候,我们需要restore variables from checkpoint files.

经常出现在checkpoint 中找不到”Tensor name not found”.

这时候需要查看一下ckpt中到底有哪些变量

  1. import os
  2. from tensorflow.python import pywrap_tensorflow
  3.  
  4. checkpoint_path = os.path.join(model_dir, "model.ckpt")
  5. # Read data from checkpoint file
  6. reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
  7. var_to_shape_map = reader.get_variable_to_shape_map()
  8. # Print tensor name and values
  9. for key in var_to_shape_map:
  10. print("tensor_name: ", key)
  11. print(reader.get_tensor(key))

可以显示ckpt中的tensor名字和值,当然也可以用pycharm调试。

补充:tensorflow中读取模型中保存的值, tf.train.NewCheckpointReader

使用tf.trian.NewCheckpointReader(model_dir)

一个标准的模型文件有一下文件, model_dir就是MyModel(没有后缀)

  1. checkpoint
  2. Model.meta
  3. Model.data-00000-of-00001
  4. Model.index
  1. import tensorflow as tf
  2. import pprint # 使用pprint 提高打印的可读性
  3. NewCheck =tf.train.NewCheckpointReader("model")

打印模型中的所有变量

  1. print("debug_string:\n")
  2. pprint.pprint(NewCheck.debug_string().decode("utf-8"))

Tensorflow 读取ckpt文件中的tensor操作

其中有3个字段, 分别是名字, 数据类型, shape

获取变量中的值

  1. print("get_tensor:\n")
  2. pprint.pprint(NewCheck.get_tensor("D/conv2d/bias"))

Tensorflow 读取ckpt文件中的tensor操作

  1. print("get_variable_to_dtype_map\n")
  2. pprint.pprint(NewCheck.get_variable_to_dtype_map())
  3. print("get_variable_to_shape_map\n")
  4. pprint.pprint(NewCheck.get_variable_to_shape_map())

以上为个人经验,希望能给大家一个参考,也希望大家多多支持我们。如有错误或未考虑完全的地方,望不吝赐教。

原文链接:https://blog.csdn.net/noirblack/article/details/71430790