如何计算从pb文件加载的tensorflow模型的flops

我想我找到了问题的原因和解决办法。下面的代码可以打印给定pb文件的触发器。

import os
import tensorflow as tf
from tensorflow.core.framework import graph_pb2
from tensorflow.python.framework import importer

os.environ['CUDA_VISIBLE_DEVICES'] = '0'

pb_path = 'mymodel.pb'

run_meta = tf.RunMetadata()
with tf.Graph().as_default():
    output_graph_def = graph_pb2.GraphDef()
    with open(pb_path, "rb") as f:
        output_graph_def.ParseFromString(f.read())
        _ = importer.import_graph_def(output_graph_def, name="")
        print('model loaded!')
    all_keys = sorted([n.name for n in tf.get_default_graph().as_graph_def().node])
    # for k in all_keys:
    #   print(k)

    with tf.Session() as sess:
        flops = tf.profiler.profile(tf.get_default_graph(), run_meta=run_meta,
            options=tf.profiler.ProfileOptionBuilder.float_operation())
        print("test flops:{:,}".format(flops.total_float_ops))

问题中打印的flops仅为18的原因是,在生成pb文件时,我将输入图像形状设置为 [None, None, 3] . 如果我改成 [500, 500, 3] ,则打印的触发器将是正确的。

在不知道输入和输出的情况下,不确定它将如何计算任何性能度量:也许它需要 CallableOptions trace_next_step and a Session 而不是手动计算。