Free open-source package to profile torch.nn.Module
modules and obtain useful information to design a model that fits your needs and constraints at development time.
With moduleprofiler
you can:
- Calculate the number of parameters of your model.
- Trace the input and output sizes of each component of your model.
- Estimate the number of operations your model performs in a forward pass.
- Calculate per module and total inference time.
All results can be obtained in one of the following formats:
dict
(default output format)pandas.DataFrame
(to perform further calculations or filtering in your code)html
(to export as webpage)LaTeX
(to include in your publications)
[Online documentation | Tutorial ]
moduleprofiler
can be installed as any regular python
module within your environment.
Install from PyPI:
python -m pip install moduleprofiler
Install from this repository:
python -m pip install git+https://github.com/eagomez2/moduleprofiler.git
You can access the online documentation. There you will find a more in depth introduction to moduleprofiler
, including tutorials, methods documentation and an extensive reference about the calculations utilized to estimate the operations of different supported torch.nn.Module
modules.
You can also run the documentation locally by going to the root folder of the package and running:
mkdocs serve
Before running this, make sure that your python environment is enabled.
If this package contributed to your work, please consider citing it:
@misc{moduleprofiler,
author = {Esteban Gómez},
title = {moduleprofiler},
year = 2024,
url = {https://github.com/eagomez2/moduleprofiler}
}
This package was developed by Esteban Gómez, member of the Speech Interaction Technology group from Aalto University.
For further details about the license of this package, please see LICENSE.