如何计算 CNN 中网络的 FLOPs [已关闭]

Question

我想设计一个占用GPU资源不超过Alexnet的卷积神经网络。我想用 FLOPs 来衡量它，但我不知道如何计算它。有没有什么工具可以做到？

Answer 1

如果您使用 Keras 和 TensorFlow 作为后端，那么您可以尝试以下示例。它计算 MobileNet 的 FLOP。

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

Answer 2

如果您使用的是 TensorFlow v1.x，则

Tobias Scheck 的答案有效，但如果您使用的是 TensorFlow v2.x，则可以使用以下代码：

import tensorflow as tf

def get_flops(model_h5_path):
    session = tf.compat.v1.Session()
    graph = tf.compat.v1.get_default_graph()
        

    with graph.as_default():
        with session.as_default():
            model = tf.keras.models.load_model(model_h5_path)

            run_meta = tf.compat.v1.RunMetadata()
            opts = tf.compat.v1.profiler.ProfileOptionBuilder.float_operation()
        
            # We use the Keras session graph in the call to the profiler.
            flops = tf.compat.v1.profiler.profile(graph=graph,
                                                  run_meta=run_meta, cmd='op', options=opts)
        
            return flops.total_float_ops

上述函数采用h5格式保存的模型的路径。您可以保存模型并以这种方式使用该功能：

model.save('path_to_my_model.h5')
tf.compat.v1.reset_default_graph()
print(get_flops('path_to_my_model.h5'))

请注意，我们使用

tf.compat.v1.reset_default_graph()

是为了避免每次调用该函数时都累积 FLOPS。