A Python to count flops of pytorch models.
Project description
Flopper - A FLOP counter for PyTorch
An FLOP counter based on fvcore with a more extensive support for any (we're trying) PyTorch modules. This tool is a lightweight wrapper around fvcore flop counter, which does all the work under the hood. We provide an easy to use API to count the number of FLOPs of any PyTorch model.
It's going to be a bit slower than fvcore, but more accurate.
Installation
pip install flopper
Usage
The simplest way to use flopper is to use the count_flops
function. It takes a model and an input batch as input and prints the total number of FLOPs.
from flopper import count_flops
model = YourRandomModel()
batch = torch.randn(1, 3, 224, 224)
flops = count_flops(model, batch)
To get more detailed information, you can do the following:
print(flops.by_operator())
print(flops.by_module())
print(flops.by_module_and_operator())
print(flops.get_table())
Adding support for custom new modules
Work in progress
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
flopper-0.1.3.tar.gz
(10.0 kB
view hashes)