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 Distribution
thop-0.0.30-1908281655.tar.gz
(4.5 kB
view hashes)
Built Distribution
Close
Hashes for thop-0.0.30.post1908281655-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 21cdc14247ce09ca53afe85c0f5a24c3145ab015f9e6ca0e7843be1f001cd6c8 |
|
MD5 | 8ebe9b3ae273d762e0a5a1b43fe0ab99 |
|
BLAKE2b-256 | 387a17fbd8c3adf2c8c0494bc3a13c0566865ea5da55bf071258036dd39fc7d0 |