神经网络参数计算--直接解析CKPT文件读取

最新推荐文章于 2024-07-05 22:53:13 发布

转载最新推荐文章于 2024-07-05 22:53:13 发布 · 575 阅读

3 ·

CC 4.0 BY-SA版权

原文链接：http://www.cnblogs.com/xiaoniu-666/p/11102818.html

文章标签：

#人工智能 #python

本文介绍了如何使用TensorFlow读取模型的ckpt文件参数，并详细展示了如何遍历和打印所有变量及其形状，以及如何计算整个网络模型的大小。

部署运行你感兴趣的模型镜像

1、tensorflow的模型文件ckpt参数获取

import tensoflow as tf
from tensorflow.python import pywrap_tensorflow
 
model_dir = "./ckpt/"
 
ckpt = tf.train.get_checkpoint_state(model_dir)
ckpt_path = ckpt.model_checkpoint_path
 
reader = pywrap_tensorflow.NewCheckpointReader(ckpt_path)
param_dict = reader.get_variable_to_shape_map()
 
for key, val in param_dict.items():
    try:
        print key, val
    except:

2、参数计算（求网络模型大小）

from tensorflow.python import pywrap_tensorflow
import os
import numpy as np
model_dir = "models_pretrained/"
checkpoint_path = os.path.join(model_dir, "model.ckpt-82798")
reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
var_to_shape_map = reader.get_variable_to_shape_map()
total_parameters = 0
for key in var_to_shape_map:#list the keys of the model
    # print(key)
    # print(reader.get_tensor(key))
    shape = np.shape(reader.get_tensor(key))  #get the shape of the tensor in the model
    shape = list(shape)
    # print(shape)
    # print(len(shape))
    variable_parameters = 1
    for dim in shape:
        # print(dim)
        variable_parameters *= dim
    # print(variable_parameters)
    total_parameters += variable_parameters

print(total_parameters)