TensorFlow:有没有办法测量模型的FLOPS?

时间:2017-07-13 16:05:47

标签: python tensorflow

我能找到的最接近的例子是:https://github.com/tensorflow/tensorflow/issues/899

使用此最小可重现代码:

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文件?

4 个答案:

答案 0 :(得分:12)

我想在Tobias Schnek的回答的基础上回答以及回答原始问题:如何从pb文件中获取FLOP。

运行Tobias的第一段代码,回答TensorFlow 1.6.0

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)

    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('TF stats gives',flops.total_float_ops)

我们得到以下输出:

Flops should be ~ 7200
TF stats gives 8288

那么,为什么我们得到8288而不是期望的结果7200=2*25*16*9 [a] ?答案与张量AB初始化的方式相同。使用高斯分布进行初始化会花费一些FLOP。通过

更改AB的定义
    A = tf.Variable(initial_value=tf.zeros([25, 16]))
    B = tf.Variable(initial_value=tf.zeros([16, 9]))

给出预期输出7200

通常,网络的变量在其他方案中使用高斯分布进行初始化。大多数时候,我们对初始化FLOP不感兴趣,因为它们在初始化期间完成一次,并且在训练期间也不会发生,也不会发生推断。那么,如何才能获得FLOP的确切数量而忽略初始化FLOP

使用pb冻结图表。从pb文件计算FLOP实际上是OP的用例。

以下代码段说明了这一点:

import tensorflow as tf
from tensorflow.python.framework import graph_util

def load_pb(pb):
    with tf.gfile.GFile(pb, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name='')
        return graph

# ***** (1) Create Graph *****
g = tf.Graph()
sess = tf.Session(graph=g)
with g.as_default():
    A = tf.Variable(initial_value=tf.random_normal([25, 16]))
    B = tf.Variable(initial_value=tf.random_normal([16, 9]))
    C = tf.matmul(A, B, name='output')
    sess.run(tf.global_variables_initializer())
    flops = tf.profiler.profile(g, options = tf.profiler.ProfileOptionBuilder.float_operation())
    print('FLOP before freezing', flops.total_float_ops)
# *****************************        

# ***** (2) freeze graph *****
output_graph_def = graph_util.convert_variables_to_constants(sess, g.as_graph_def(), ['output'])

with tf.gfile.GFile('graph.pb', "wb") as f:
    f.write(output_graph_def.SerializeToString())
# *****************************


# ***** (3) Load frozen graph *****
g2 = load_pb('./graph.pb')
with g2.as_default():
    flops = tf.profiler.profile(g2, options = tf.profiler.ProfileOptionBuilder.float_operation())
    print('FLOP after freezing', flops.total_float_ops)

输出

FLOP before freezing 8288
FLOP after freezing 7200

[a] 对于产品AB而言,矩阵乘法的FLO​​P通常为mq(2p -1),其中A[m, p]B[p, q]但TensorFlow由于某种原因返回2mpq 。已打开issue以了解原因。

答案 1 :(得分:9)

有点晚了,但未来可能会对一些访客有所帮助。对于您的示例,我成功测试了以下代码段:

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结合使用,如下面的代码段所示:

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))

我希望我能帮忙!

答案 2 :(得分:0)

由于不推荐使用探查器方法并将其移至compat.v1下,因此上述方法不再适用于TF2.0。似乎仍然需要实现此功能。

以下是有关Github的问题: https://github.com/tensorflow/tensorflow/issues/32809

答案 3 :(得分:0)

Another user 发布了一个答案。它被mod删除,因此无法恢复。但它确实解决了问题,并且比其他答案更好。所以我在这里重复一遍。


您可以使用以下 pip 包来获取一些基本信息,例如模型的内存要求,否。参数、触发器等

https://pypi.org/project/model-profiler

它会输出类似的东西

<头>
模型简介 价值 单位
选定的 GPU ['0', '1'] GPU ID
没有。 FLOPs 0.30932349055999997 BFLOPs
GPU 内存要求 7.4066760912537575 GB
模型参数 138.357544 百万
模型权重所需的内存 527.7921447753906 MB

用法

[从图书馆网站逐字复制]

from tensorflow.keras.applications import VGG16

model = VGG16(include_top=True)

from model_profiler import model_profiler

Batch_size = 128
profile = model_profiler(model, Batch_size)

print(profile)