Keras/Pytorch neural network size, operations and parameters counter. Relies on tensorflow profiler for tf/keras models, torchinfo (https://github.com/TylerYep/torchinfo) for torch models, and tflite tools (https://github.com/eliberis/tflite-tools) for tensorflow lite model profiling.

For tf/Keras models, the input size gets inferred from the first layer of the network, so avoid [None] undefined dimensions. Batch size gets automatically set to 1.

[Keras]:

import nessi

nessi.get_model_size(model, 'keras')

For torch models, you need to specify the input size as a touple when calling get_model_size.

[Torch]:

import nessi

nessi.get_model_size(net, 'torch' ,input_size=(1,3,320,320))

For tflite models, either pass the model path or the bytearray to get_model_size.

[TFLite]:

import nessi

nessi.get_model_size(tflite_model OR path/to/model.tflite, 'tflite')