tensorflow tf flops计算

https://pypi.org/project/keras-flops/

from tensorflow.keras import Model, Input
from tensorflow.keras.layers import Dense, Flatten, Conv2D, MaxPooling2D, Dropout

from keras_flops import get_flops
# build model

# build model
# inp = Input((32, 32, 1))
# x = Conv2D(32, kernel_size=(3, 3), activation="relu")(inp)
# x = MaxPooling2D(pool_size=(2, 2))(x)
# x = Flatten()(x)
# x = Dense(128, activation="relu")(x)
# out = Dense(10, activation="softmax")(x)
# model = Model(inp, out)

# # build model
# build model
inp = Input((32, 32, 1))
x = Conv2D(32, 3, padding='same', activation='relu')(inp)
#x = Conv2D(64, (3, 3), activation="relu")(x)
x = MaxPooling2D(pool_size=(2, 2))(x)
x = Dropout(0.25)(x)
x = Flatten()(x)
x = Dense(128, activation="relu")(x)
x = Dropout(0.5)(x)
out = Dense(10, activation="softmax")(x)
model = Model(inp, out)
model.summary()
# calculae flops
flops = get_flops(model, batch_size=1)
print(f"flops: {flops / 10 ** 9:.03} g")

结果
_TFProfRoot (–/2.76m flops)
functional_1/dense/MatMul (2.10m/2.10m flops)
functional_1/conv2d/Conv2D (589.82k/589.82k flops)
functional_1/conv2d/BiasAdd (32.77k/32.77k flops)
functional_1/max_pooling2d/MaxPool (32.77k/32.77k flops)
functional_1/dense_1/MatMul (2.56k/2.56k flops)
functional_1/dense/BiasAdd (128/128 flops)
functional_1/dense_1/Softmax (50/50 flops)
functional_1/dense_1/BiasAdd (10/10 flops)
End of Report====
flops: 0.00276 g