tensorflow - 如何计算用从 .pb 文件加载的图形定义的 tensorflow 模型中可训练参数的总数？

Question

我想计算 tensorflow 模型中的参数。它与现有问题类似，如下所示。

How to count total number of trainable parameters in a tensorflow model?

但如果模型是用从 .pb 文件加载的图形定义的，则所有建议的答案都不起作用。基本上我用以下函数加载了图表。

def load_graph(model_file):

  graph = tf.Graph()
  graph_def = tf.GraphDef()

  with open(model_file, "rb") as f:
    graph_def.ParseFromString(f.read())

  with graph.as_default():
    tf.import_graph_def(graph_def)

  return graph

一个示例是在 tensorflow-for-poets-2 中加载 frozen_graph.pb 文件以进行再训练。

https://github.com/googlecodelabs/tensorflow-for-poets-2

Answer 1

据我了解，GraphDef 没有足够的信息来描述Variables。正如解释的那样 here ，您将需要 MetaGraph，其中包含 GraphDef 和 CollectionDef，后者是可以描述 Variables 的映射。所以下面的代码应该给我们正确的可训练变量计数。

导出元图:

import tensorflow as tf

a = tf.get_variable('a', shape=[1])
b = tf.get_variable('b', shape=[1], trainable=False)
init = tf.global_variables_initializer()
saver = tf.train.Saver([a])

with tf.Session() as sess:
    sess.run(init)
    saver.save(sess, r'.\test')

导入元图并统计可训练参数的总数。

import tensorflow as tf

saver = tf.train.import_meta_graph('test.meta')

with tf.Session() as sess:
    saver.restore(sess, 'test')

total_parameters = 0
for variable in tf.trainable_variables():
    total_parameters += 1
print(total_parameters)