TensorFlow2.0计算flops

可以使用TensorFlow2.0内置的tf.profiler模块来计算FLOPS。以下是一个简单的示例代码：

import tensorflow as tf

# 构建模型
model = tf.keras.models.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 生成输入数据
input_data = tf.random.normal([1, 28, 28, 1])

# 运行一次前向传播以生成计算图
@tf.function
def inference():
    output = model(input_data)
    return output

output = inference()

# 打印FLOPS
tf.profiler.profile(
    tf.get_default_graph(),
    options=tf.profiler.ProfileOptionBuilder.float_operation())

执行上述代码后，会输出模型的总FLOPS数量。注意，这里的FLOPS是指在单次前向传播中所执行的总乘法和加法操作数量，而不是训练过程中的总操作数量。因此，如果您的模型中包含了dropout、batch normalization等操作，计算得到的FLOPS数量会偏低。