使用Tensorflow怎么對ckpt文件中的tensor進行讀取

本篇文章給大家分享的是有關使用Tensorflow怎么對ckpt文件中的tensor進行讀取，小編覺得挺實用的，因此分享給大家學習，希望大家閱讀完這篇文章后可以有所收獲，話不多說，跟著小編一起來看看吧。

在使用pre-train model時候，我們需要restore variables from checkpoint files.

經常出現在checkpoint 中找不到”Tensor name not found”.

這時候需要查看一下ckpt中到底有哪些變量

import os
from tensorflow.python import pywrap_tensorflow

checkpoint_path = os.path.join(model_dir, "model.ckpt")
# Read data from checkpoint file
reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
# Print tensor name and values
for key in var_to_shape_map:
  print("tensor_name: ", key)
  print(reader.get_tensor(key))

可以顯示ckpt中的tensor名字和值，當然也可以用pycharm調試。

補充：tensorflow中讀取模型中保存的值, tf.train.NewCheckpointReader

使用tf.trian.NewCheckpointReader(model_dir)

一個標準的模型文件有一下文件, model_dir就是MyModel(沒有后綴)

checkpoint
Model.meta
Model.data-00000-of-00001
Model.index

import tensorflow as tf
import pprint # 使用pprint 提高打印的可讀性
NewCheck =tf.train.NewCheckpointReader("model")

打印模型中的所有變量

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

其中有3個字段, 分別是名字, 數據類型, shape

獲取變量中的值

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

print("get_variable_to_dtype_map\n")
pprint.pprint(NewCheck.get_variable_to_dtype_map())
print("get_variable_to_shape_map\n")
pprint.pprint(NewCheck.get_variable_to_shape_map())

以上就是使用Tensorflow怎么對ckpt文件中的tensor進行讀取，小編相信有部分知識點可能是我們日常工作會見到或用到的。希望你能通過這篇文章學到更多知識。更多詳情敬請關注億速云行業資訊頻道。