TensorFlow:有没有办法测量模型的 FLOPS? [英] TensorFlow: Is there a way to measure FLOPS for a model?

问题描述

我能得到的最接近的例子在这个问题中找到:https://github.com/张量流/张量流/问题/899

The closest example I can get is found in this issue: https://github.com/tensorflow/tensorflow/issues/899

使用这个最小的可重现代码:

With this minimum reproducible code:

import tensorflow as tf
import tensorflow.python.framework.ops as ops 
g = tf.Graph()
with g.as_default():
  A = tf.Variable(tf.random_normal( [25,16] ))
  B = tf.Variable(tf.random_normal( [16,9] ))
  C = tf.matmul(A,B) # shape=[25,9]
for op in g.get_operations():
  flops = ops.get_stats_for_node_def(g, op.node_def, 'flops').value
  if flops is not None:
    print 'Flops should be ~',2*25*16*9
    print '25 x 25 x 9 would be',2*25*25*9 # ignores internal dim, repeats first
    print 'TF stats gives',flops

但是，返回的 FLOPS 始终为 None.有没有办法具体测量FLOPS，尤其是PB文件?

However, the FLOPS returned is always None. Is there a way to concretely measure FLOPS, especially with a PB file?

推荐答案

有点晚了，但也许对未来的一些访问者有所帮助.对于您的示例，我成功测试了以下代码段:

A little bit late but maybe it helps some visitors in future. For your example I successfully tested the following snippet:

g = tf.Graph()
run_meta = tf.RunMetadata()
with g.as_default():
    A = tf.Variable(tf.random_normal( [25,16] ))
    B = tf.Variable(tf.random_normal( [16,9] ))
    C = tf.matmul(A,B) # shape=[25,9]

    opts = tf.profiler.ProfileOptionBuilder.float_operation()    
    flops = tf.profiler.profile(g, run_meta=run_meta, cmd='op', options=opts)
    if flops is not None:
        print('Flops should be ~',2*25*16*9)
        print('25 x 25 x 9 would be',2*25*25*9) # ignores internal dim, repeats first
        print('TF stats gives',flops.total_float_ops)

也可以将分析器与 Keras 结合使用，如下所示:

It's also possible to use the profiler in combination with Keras like the following snippet:

import tensorflow as tf
import keras.backend as K
from keras.applications.mobilenet import MobileNet

run_meta = tf.RunMetadata()
with tf.Session(graph=tf.Graph()) as sess:
    K.set_session(sess)
    net = MobileNet(alpha=.75, input_tensor=tf.placeholder('float32', shape=(1,32,32,3)))

    opts = tf.profiler.ProfileOptionBuilder.float_operation()    
    flops = tf.profiler.profile(sess.graph, run_meta=run_meta, cmd='op', options=opts)

    opts = tf.profiler.ProfileOptionBuilder.trainable_variables_parameter()    
    params = tf.profiler.profile(sess.graph, run_meta=run_meta, cmd='op', options=opts)

    print("{:,} --- {:,}".format(flops.total_float_ops, params.total_parameters))

希望能帮到你！

这篇关于TensorFlow:有没有办法测量模型的 FLOPS?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！