2

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

如何计算张量流模型中可训练参数的总数？

但是，如果模型是使用从 .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

4

1 回答 1

1

据我了解， aGraphDef没有足够的信息来描述Variables。正如这里所解释的，您将需要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)
于 2018-05-03T20:26:01.343 回答