TensorFlow：有没有一种方法可以测量模型的FLOPS？

Question

我能提供的最接近的例子在这个问题中找到：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 文件？

- kwotsin

---

在TF 2.0上如何实现？ - Eduardo Reis

你可以在TF2中使用model_profiler。 - Safi

Answer 1

我希望在托比亚斯·施内克的回答基础上进行补充，同时回答原始问题：如何从pb文件中获取FLOP。

在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]？答案在于张量A和B的初始化方式。用高斯分布进行初始化会消耗一些FLOP。通过改变A和B的定义来解决：

    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] 矩阵乘法的 FLOP 通常是 mq(2p-1) 其中 A[m,p] 和 B[p,q] 的乘积 AB，但是 TensorFlow 出于某些原因返回了 2mpq。已经开放了一个问题以了解原因。

---

- BiBi

---

冻结图确实有效！！谢谢您的回答！！ - Dheeraj Peri

你好，感谢提供详细信息！我有一个更多的问题需要解决：我收到了带有二维输入张量的冻结图形（.pb），其中形状没有完全定义（例如[？, 3]），因此您的代码给出了FLOP after freezing 0 12 ops no flops stats due to incomplete shapes.有没有一种方法可以只是“假设”一个特定的形状而不必重新进行培训过程呢？ - tre95

Answer 2

可能略有迟缓，但希望对于未来的一些访问者能有所帮助。针对您的示例，我成功地测试了以下片段：

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

希望我能帮上忙！

- Tobias Scheck

---

第一个片段输出“Flops应该是大约7200”和“TF统计值为8288”。为什么会有这样的差异？我在此基础上解释一下。 - BiBi

它运行得非常好。但是有没有办法防止 tf.profiler.profile 记录到控制台？ - Zz Tux

Answer 3

由于TF2.0中已经取消了分析器方法并移至compat.v1，因此上述方法不再适用。看起来这个功能仍需要被实现。

以下是Github上的一个问题: https://github.com/tensorflow/tensorflow/issues/32809

- Madhavan Seshadri

---

这个版本是什么情况？https://gist.github.com/scheckmedia/cadc5eb3d74ed57a4f3d78011a9f6f7c - Tobias Scheck

Answer 4

另一个用户发布了一个回答，由于被管理员删除，因此无法恢复。但是它确实解决了问题，并且比其他答案更好。所以我在这里重复一遍。

---

您可以使用以下pip包获取一些基本信息，例如模型的内存需求，参数数量，FLOPS等。

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)

- mafu

---

注意：从源代码来看，这个库似乎没有实现所有常见的层。因此输出可能会非常不准确。由于该库还很年轻，作者将来可能会修复这个问题。 - mafu