A tool to layer-wise count the MACs and parameters of PyTorch model.
Project description
PyTorch-layerwise-OpCounter
A tool for profile the MACs, parameters, input_shape, output_shape et.al of each layer in Pytorch model. Forked from Lyken17/pytorch-OpCounter which is not supporting layer-wise profile and I will follow it.
How to install
pip install torchlop
OR
pip install --upgrade git+https://github.com/hahnyuan/pytorch-layerwise-OpCounter.git
How to use
-
Basic usage
from torchvision.models import resnet50 from torchlop import profile model = resnet50() input = torch.randn(1, 3, 224, 224) macs, params, layer_infos = profile(model, inputs=(input, ))
-
The layer_infos is a dict that contains the infos for each layer of the pytorch model.
- The key is the name of the layer
- 'type': the class name of the layer
- 'in_size': input size
- 'out_size': output size
- 'ops': operations (MACs)
- 'params': parameters
-
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) macs, params, ayer_infos = profile(model, inputs=(input, ), custom_ops={YourModule: count_your_model})
Results of Recent Models
The implementation are adapted from torchvision
. Following results can be obtained using benchmark/evaluate_famours_models.py.
|
|
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.