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.
pip install torchlop
OR
pip install --upgrade git+https://github.com/hahnyuan/pytorch-layerwise-OpCounter.git
-
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, layer_infos = profile(model, inputs=(input, ), custom_ops={YourModule: count_your_model})
-
Write the layerwise profile information into a csv file
from torchvision.models import resnet50 from torchlop import profile from torchlop.rst_process import write_csv model = resnet50() input = torch.randn(1, 3, 224, 224) macs, params, layer_infos = profile(model, inputs=(input, )) csv_file='profile.csv' write_csv(csv_file,layer_infos)
The implementation are adapted from torchvision
. Following results can be obtained using benchmark/evaluate_famours_models.py.
|
|