Tiny PyTorch helper package to get stats on GPU usage, model params, etc.
Project description
flaregun
A simple helper package for PyTorch to get better visibility on real-time GPU usage, model params, etc.
Installation
$ pip install flaregun
Usage
Get real-time Nvidia GPU memory usage in Python:
from flaregun import GPUStats
# Pretty print statistics for GPU #0
GPUStats(device=0).print()
> "GPU memory usage: 3061 / 32510 MB"
Get parameter count in any PyTorch compatible model (e.g. HuggingFace, etc.):
from flaregun import ModelStats
# Get HuggingFace model
model = AutoModelForMaskedLM.from_pretrained("/path/to/Longformer-Model")
# Pretty print Model parameter count
ModelStats(model).print()
> "148711257 params (148711257 trainable | 0 non-trainable)"
API
All features of the library are listed below.
GPU memory utilization
from flaregun import GPUStats
device = ...integer of GPU device...
# Free GPU memory (in MB)
free_mem = GPUStats(device).free()
# Total GPU memory (in MB)
total_mem = GPUStats(device).total()
# Used GPU memory (in MB)
used_mem = GPUStats(device).used()
Model parameter count
from flaregun import ModelStats
model = ...PyTorch-compatible model...
# Total params
total = ModelStats(model).total()
# Trainable params
trainable = ModelStats(model).trainable()
# Frozen (non-trainable) params
frozen = ModelStats(model).frozen()
Contributing
Interested in contributing? Check out the contributing guidelines. Please note that this project is released with a Code of Conduct. By contributing to this project, you agree to abide by its terms.
Running Tests
poetry run pytest tests/
Test PyPI
Create release:
git tag vXX.XX.XXX
git push --tags
# Then go to Github and link this tag to a Release
Publish:
poetry build
poetry publish -r test-pypi --username XXX --password XXXX
Install:
pip3 install --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple flaregun
PyPI
Publish:
poetry publish
Building Docs
cd docs
poetry run make clean html && poetry run make html
License
flaregun
was created by Michael Wornow. It is licensed under the terms of the MIT license.
Credits
flaregun
was created with cookiecutter
and the py-pkgs-cookiecutter
template.
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.