当前位置 博文首页 > Tensorflow 读取ckpt文件中的tensor操作

    Tensorflow 读取ckpt文件中的tensor操作

    作者:noirblack 时间:2021-07-17 18:52

    在使用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())

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

    jsjbwy
    下一篇:没有了