keras-flops 0.1.2

Description:

kerasflops 0.1.2

keras-flops


FLOPs calculator for neural network architecture written in tensorflow (tf.keras) v2.2+
This stands on the shoulders of giants, tf.profiler.
Requirements

Python 3.6+
Tensorflow 2.2+

Installation
Using pip:
pip install keras-flops

Example
See colab examples here in details.
from tensorflow.keras import Model, Input
from tensorflow.keras.layers import Dense, Flatten, Conv2D, MaxPooling2D, Dropout

from keras_flops import get_flops

# build model
inp = Input((32, 32, 3))
x = Conv2D(32, kernel_size=(3, 3), 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)

# Calculae FLOPS
flops = get_flops(model, batch_size=1)
print(f"FLOPS: {flops / 10 ** 9:.03} G")
# >>> FLOPS: 0.0338 G

Support
Support tf.keras.layers as follows,



name
layer




Conv
Conv[1D/2D/3D]



Conv[1D/2D]Transpose



DepthwiseConv2D



SeparableConv[1D/2D]


Pooling
AveragePooling[1D/2D]



GlobalAveragePooling[1D/2D/3D]



MaxPooling[1D/2D]



GlobalMaxPool[1D/2D/3D]


Normalization
BatchNormalization


Activation
Softmax


Attention
Attention



AdditiveAttention


others
Dense



Not supported
Not support tf.keras.layers as follows. They are calculated as zero or smaller value than correct value.



name
layer




Conv
Conv3DTranspose


Pooling
AveragePooling3D



MaxPooling3D



UpSampling[1D/2D/3D]


Normalization
LayerNormalization


RNN
SimpleRNN



LSTM



GRU


others
Embedding