A tool to count the FLOPs of PyTorch model.
Project description
THOP: PyTorch-OpCounter
How to install
-
Through PyPi
pip install thop
-
Using GitHub (always latest)
pip install --upgrade git+https://github.com/Lyken17/pytorch-OpCounter.git
How to use
-
Basic usage
from torchvision.models import resnet50 from thop import profile model = resnet50() input = torch.randn(1, 3, 224, 224) flops, params = profile(model, inputs=(input, ))
-
Define the rule for 3rd party module.
class YourModule(nn.Module): # your definition def count_your_model(model, x, y): # your rule here input = torch.randn(1, 3, 224, 224) flops, params = profile(model, inputs=(input, ), custom_ops={YourModule: count_your_model})
-
Improve the output readability
Call
thop.clever_format
to give a better format of the output.from thop import clever_format flops, params = clever_format([flops, params], "%.3f")
Results on Recent Models
|
|
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 Distributions
No source distribution files available for this release.See tutorial on generating distribution archives.
Built Distribution
Close
Hashes for thop-0.0.30.post1908282106-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | afcf5ea5f53835c1c338b514322b8ea0834fab759616aff408e82095e68ba7c2 |
|
MD5 | a728a3b5c16b750b3652ffb432149915 |
|
BLAKE2b-256 | a873b06e9cbed394b95b9066e61bb7106fccf561758ebfe466665de4837f03ac |