适用于TensorFlow2.0,计算flops

在TensorFlow 2.0中，可以使用tf.profiler包来计算模型的FLOPS。以下是计算FLOPS的示例代码：

import tensorflow as tf

# 构建模型
model = tf.keras.applications.ResNet50()

# 输入形状
input_shape = (1, 224, 224, 3)

# 定义输入
inputs = tf.keras.Input(shape=input_shape[1:])

# 调用模型
outputs = model(inputs)

# 构建模型
model = tf.keras.Model(inputs=inputs, outputs=outputs)

# 运行一次模型以使TensorFlow分配所有变量
_ = model(tf.zeros(input_shape))

# 计算FLOPS
flops = tf.profiler.profile(
    tf.get_default_graph(),
    options=tf.profiler.ProfileOptionBuilder.float_operation())
print('FLOPS:', flops.total_float_ops)

在此示例中，我们使用ResNet50作为模型，并定义输入形状。然后，我们通过调用模型来构建完整的模型。接下来，我们运行一次模型以使TensorFlow分配所有变量。最后，我们使用`tf.profiler.profile()`函数计算模型的FLOPS，并将其打印出来。