A simple helper package for PyTorch to get better visibility on real-time GPU usage, model params, etc.
from flaregun import GPUStats, ModelStats
GPUStats(device=0).print()
> "GPU memory usage: 3061 / 32510 MB"
from transformers import AutoModel
ModelStats(AutoModel.from_pretrained('gpt2')).print()
> "124439808 params (124439808 trainable | 0 non-trainable)"
$ pip install flaregun
Get real-time Nvidia GPU memory usage in Python:
from flaregun import GPUStats
# Pretty print statistics for Nvidia 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
from transformers import AutoModel
model = AutoModel.from_pretrained("gpt2")
# Pretty print Model parameter count
ModelStats(model).print()
> "124439808 params (124439808 trainable | 0 non-trainable)"
All features of the library are listed below.
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()
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()
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.
poetry run pytest tests/
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
Publish:
poetry publish --username XXX --password XXXX
cd docs
poetry run make clean html && poetry run make html
flaregun
was created by Michael Wornow. It is licensed under the terms of the MIT license.
flaregun
was created with cookiecutter
and the py-pkgs-cookiecutter
template.