Skip to content

FLOPs calculator with tf.profiler for neural network architecture written in tensorflow 2.2+ (tf.keras)

License

Notifications You must be signed in to change notification settings

tokusumi/keras-flops

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

15 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

keras-flops

PyPI version

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

About

FLOPs calculator with tf.profiler for neural network architecture written in tensorflow 2.2+ (tf.keras)

Topics

Resources

License

Stars

Watchers

Forks

Sponsor this project

 

Packages

No packages published