Estimate FLOPs of neural networks
Project description
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:
pip install pthflops
or directly from the github repository:
git clone https://github.com/1adrianb/pytorch-estimate-flops && pytorch-estimate-flops
python setup.py install
Example
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)
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
pthflops-0.1.0.tar.gz
(3.7 kB
view hashes)
Built Distribution
pthflops-0.1.0-py3.6.egg
(6.3 kB
view hashes)