CODEWITHSUNDEEP
pytorchflops
afsara_ben
afsara_ben

Reputation: 682

calculate flops in a custom pytorch model

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?

Upvotes: 2

Views: 13438

Answers (2)

afsara_ben
afsara_ben

Reputation: 682

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.

Upvotes: 2

yeachan park
yeachan park

Reputation: 180

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__").

Upvotes: 1

Related Questions