FLOPs calculator with tf.profiler for neural network architecture written in tensorflow 2.2+ (tf.keras)
Project description
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 |
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
keras-flops-0.1.2.tar.gz
(4.4 kB
view hashes)
Built Distribution
Close
Hashes for keras_flops-0.1.2-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 1202bb610f82a931ed65964fb9c330875001a5b1d527b279062185178541366a |
|
MD5 | 5eb153df9646dd94bba5be434624449a |
|
BLAKE2b-256 | bc77e384f6e427e1920624a9229b411010eaa244134cff6a230d303d3b1ea0be |