获得tensorflow模型的参数量

CHNguoshiwushuang

于 2018-08-11 17:25:58 发布

阅读量1w

点赞数 16

CC 4.0 BY-SA版权

分类专栏： tensorflow笔记

本文链接：https://blog.csdn.net/CHNguoshiwushuang/article/details/81588794

tensorflow模型训练好之后，通常会保存为.ckpt文件，有时我们想了解一下模型保存了多少参数，这固然可以手动计算，但是速度太慢，这里我写了一个程序可以直接获得模型的参数量，下面是源代码：

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(re