https://github.com/zhijian-liu/torchprofile
A general and accurate MACs / FLOPs profiler for PyTorch models
https://github.com/zhijian-liu/torchprofile
profiler pytorch
Last synced: 5 months ago
JSON representation
A general and accurate MACs / FLOPs profiler for PyTorch models
- Host: GitHub
- URL: https://github.com/zhijian-liu/torchprofile
- Owner: zhijian-liu
- License: mit
- Created: 2019-09-03T00:30:33.000Z (about 6 years ago)
- Default Branch: master
- Last Pushed: 2024-05-05T00:01:21.000Z (over 1 year ago)
- Last Synced: 2025-04-13T14:09:14.042Z (6 months ago)
- Topics: profiler, pytorch
- Language: Python
- Homepage: https://pypi.org/project/torchprofile/
- Size: 79.1 KB
- Stars: 603
- Watchers: 10
- Forks: 42
- Open Issues: 7
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# Torchprofile
This is a profiler to count the number of MACs / FLOPs of PyTorch models based on `torch.jit.trace`.
* It is more **general** than ONNX-based profilers as some operations in PyTorch are not supported by ONNX for now.
* It is more **accurate** than hook-based profilers as they cannot profile operations within `torch.nn.Module`.
## Installation
```bash
pip install torchprofile
```
## Getting Started
You should first define your PyTorch model and its (dummy) input:
```python
import torch
from torchvision.models import resnet18
model = resnet18()
inputs = torch.randn(1, 3, 224, 224)
```
You can then measure the number of MACs using `profile_macs`:
```python
from torchprofile import profile_macs
macs = profile_macs(model, inputs)
```
## License
This repository is released under the MIT license. See [LICENSE](LICENSE) for additional details.