calculate flops in a custom pytorch model

Question

I have a deeply nested pytorch model and want to calculate the flops per layer. I tried using the flopth, ptflops, pytorch-OpCounter library but couldn't run it for such a deeply nested model. How to calculate the number of mul/add operations and flops each layer in this model?

Answer 1

Using Flop Counter for PyTorch Models worked.

The following was mentioned in ptflops because of which my custom model faced errors -

This script doesn't take into account torch.nn.functional.* operations. For an instance, if one have a semantic segmentation model and use torch.nn.functional.interpolate to upscale features, these operations won't contribute to overall amount of flops. To avoid that one can use torch.nn.Upsample instead of torch.nn.functional.interpolate.

Answer 2

Please use PyProf (https://github.com/NVIDIA/PyProf) package. It is well explained in the page (https://github.com/NVIDIA/PyProf/blob/main/docs/profile.rst).

If you write the code "net.py" in style described in the above page, please run the followling script:

nsys profile -f true -o net --export sqlite train net.py
python -m pyprof.parse net.sqlite > net.dict
python -m pyprof.prof --csv -c idx,dir,op,kernel,params,sil,flops,bytes net.dict > results.csv

Then you get the FLOPs in the column of result.csv file.

If you get some AssertionErrors (typically this is due to undefined operations in the PyProf package), you should manually add some operations in the ~/site-packages/pyprof/prof directory. (directory where PyProf package is installed).

For instance, if you get the AssertionError about "__iadd__", then go to pointwise.py and add "__iadd__" in the proper position. (with the same line as "__add__").