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

Question

我想计算张量流模型中的参数。它类似于现有的问题，如下所示。

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 中加载 freeze_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')

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

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)