[](https://opensource.org/licenses/BSD-3-Clause) [](https://travis-ci.com/1adrianb/pytorch-estimate-flops)
[](https://pypi.org/project/pthflops/)
# pytorch-estimate-flops
Simple pytorch utility that estimates the number of FLOPs for a given network. For now only some basic operations are supported (basically the ones I needed for my models). More will be added soon.
All contributions are welcomed.
## Installation
You can install the model using pip:
```bash
pip install pthflops
```
or directly from the github repository:
```bash
git clone https://github.com/1adrianb/pytorch-estimate-flops && cd pytorch-estimate-flops
python setup.py install
```
Note: pytorch 1.8 or newer is recommended.
## Example
```python
import torch
from torchvision.models import resnet18
from pthflops import count_ops
# Create a network and a corresponding input
device = 'cuda:0'
model = resnet18().to(device)
inp = torch.rand(1,3,224,224).to(device)
# Count the number of FLOPs
count_ops(model, inp)
```
Ignoring certain layers:
```python
import torch
from torch import nn
from pthflops import count_ops
class CustomLayer(nn.Module):
def __init__(self):
super(CustomLayer, self).__init__()
self.conv1 = nn.Conv2d(5, 5, 1, 1, 0)
# ... other layers present inside will also be ignored
def forward(self, x):
return self.conv1(x)
# Create a network and a corresponding input
inp = torch.rand(1,5,7,7)
net = nn.Sequential(
nn.Conv2d(5, 5, 1, 1, 0),
nn.ReLU(inplace=True),
CustomLayer()
)
# Count the number of FLOPs, jit mode:
count_ops(net, inp, ignore_layers=['CustomLayer'])
# Note: if you are using python 1.8 or newer with fx instead of jit, the naming convention changed. As such, you will have to pass ['_2_conv1']
# Please check your model definition to account for this.
# Count the number of FLOPs, fx mode:
count_ops(net, inp, ignore_layers=['_2_conv1'])
```
Raw data
{
"_id": null,
"home_page": "https://github.com/1adrianb/pytorch-estimate-flops",
"name": "pthflops",
"maintainer": "",
"docs_url": null,
"requires_python": "",
"maintainer_email": "",
"keywords": "",
"author": "Adrian Bulat",
"author_email": "adrian@adrianbulat.com",
"download_url": "https://files.pythonhosted.org/packages/51/e9/610f95556b06a015e70e37bbeee80f8a57e87e5571c6be6c9901551532a1/pthflops-0.4.2.tar.gz",
"platform": null,
"description": "[](https://opensource.org/licenses/BSD-3-Clause) [](https://travis-ci.com/1adrianb/pytorch-estimate-flops)\n[](https://pypi.org/project/pthflops/)\n\n# pytorch-estimate-flops\n\nSimple pytorch utility that estimates the number of FLOPs for a given network. For now only some basic operations are supported (basically the ones I needed for my models). More will be added soon.\n\nAll contributions are welcomed.\n\n## Installation\n\nYou can install the model using pip:\n\n```bash\npip install pthflops\n```\nor directly from the github repository:\n```bash\ngit clone https://github.com/1adrianb/pytorch-estimate-flops && cd pytorch-estimate-flops\npython setup.py install\n```\n\nNote: pytorch 1.8 or newer is recommended.\n\n## Example\n\n```python\nimport torch\nfrom torchvision.models import resnet18\n\nfrom pthflops import count_ops\n\n# Create a network and a corresponding input\ndevice = 'cuda:0'\nmodel = resnet18().to(device)\ninp = torch.rand(1,3,224,224).to(device)\n\n# Count the number of FLOPs\ncount_ops(model, inp)\n```\n\nIgnoring certain layers:\n\n```python\nimport torch\nfrom torch import nn\nfrom pthflops import count_ops\n\nclass CustomLayer(nn.Module):\n def __init__(self):\n super(CustomLayer, self).__init__()\n self.conv1 = nn.Conv2d(5, 5, 1, 1, 0)\n # ... other layers present inside will also be ignored\n\n def forward(self, x):\n return self.conv1(x)\n\n# Create a network and a corresponding input\ninp = torch.rand(1,5,7,7)\nnet = nn.Sequential(\n nn.Conv2d(5, 5, 1, 1, 0),\n nn.ReLU(inplace=True),\n CustomLayer()\n)\n\n# Count the number of FLOPs, jit mode:\ncount_ops(net, inp, ignore_layers=['CustomLayer'])\n\n# Note: if you are using python 1.8 or newer with fx instead of jit, the naming convention changed. As such, you will have to pass ['_2_conv1']\n# Please check your model definition to account for this.\n# Count the number of FLOPs, fx mode:\ncount_ops(net, inp, ignore_layers=['_2_conv1'])\n\n```\n\n\n",
"bugtrack_url": null,
"license": "BSD",
"summary": "Estimate FLOPs of neural networks",
"version": "0.4.2",
"project_urls": {
"Homepage": "https://github.com/1adrianb/pytorch-estimate-flops"
},
"split_keywords": [],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "a0f747983544d6bc6ae37156c1b5a2b5de5ed886bbdeeda9e1eca84b5abd55cf",
"md5": "b46ec23f9310e0c7c727027ed6fe69ba",
"sha256": "8551ca3b10538cd6250b3cee8d42ced6c7cc3b955b57909feb849b95d59c45aa"
},
"downloads": -1,
"filename": "pthflops-0.4.2-py2.py3-none-any.whl",
"has_sig": false,
"md5_digest": "b46ec23f9310e0c7c727027ed6fe69ba",
"packagetype": "bdist_wheel",
"python_version": "py2.py3",
"requires_python": null,
"size": 11115,
"upload_time": "2022-05-20T23:40:23",
"upload_time_iso_8601": "2022-05-20T23:40:23.749488Z",
"url": "https://files.pythonhosted.org/packages/a0/f7/47983544d6bc6ae37156c1b5a2b5de5ed886bbdeeda9e1eca84b5abd55cf/pthflops-0.4.2-py2.py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "0c2783eab6b0a3068944d5e52cd64de9339fce58b2c8e33fec08bafa89843881",
"md5": "298a671a1ecf4b0b7eafe04ede1ae5ca",
"sha256": "7211664c4d47032c6859a84d14d8ce3ccd17ac08028389b92c8d66bf6ca2580c"
},
"downloads": -1,
"filename": "pthflops-0.4.2-py3-none-any.whl",
"has_sig": false,
"md5_digest": "298a671a1ecf4b0b7eafe04ede1ae5ca",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": null,
"size": 11112,
"upload_time": "2022-05-20T23:40:25",
"upload_time_iso_8601": "2022-05-20T23:40:25.795320Z",
"url": "https://files.pythonhosted.org/packages/0c/27/83eab6b0a3068944d5e52cd64de9339fce58b2c8e33fec08bafa89843881/pthflops-0.4.2-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "51e9610f95556b06a015e70e37bbeee80f8a57e87e5571c6be6c9901551532a1",
"md5": "c0e226b78a267677a035bb57fc1a44d8",
"sha256": "1a64b6d75937e01cf837e3cdc688de1e0fb58a7d6105974956c3bbeaa1c105e8"
},
"downloads": -1,
"filename": "pthflops-0.4.2.tar.gz",
"has_sig": false,
"md5_digest": "c0e226b78a267677a035bb57fc1a44d8",
"packagetype": "sdist",
"python_version": "source",
"requires_python": null,
"size": 10885,
"upload_time": "2022-05-20T23:40:27",
"upload_time_iso_8601": "2022-05-20T23:40:27.534094Z",
"url": "https://files.pythonhosted.org/packages/51/e9/610f95556b06a015e70e37bbeee80f8a57e87e5571c6be6c9901551532a1/pthflops-0.4.2.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2022-05-20 23:40:27",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "1adrianb",
"github_project": "pytorch-estimate-flops",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"tox": true,
"lcname": "pthflops"
}
JSON
