torch-max-mem
This package provides decorators for memory utilization maximization with PyTorch and CUDA by starting with a maximum parameter size and applying successive halving until no more out-of-memory exception occurs.
💪
Getting Started
Assume you have a function for batched computation of nearest neighbors using brute-force distance calculation.
import torch
def knn(x, y, batch_size, k: int = 3):
return torch.cat(
[
torch.cdist(x[start : start + batch_size], y).topk(k=k, dim=1, largest=False).indices
for start in range(0, x.shape[0], batch_size)
],
dim=0,
)
With torch_max_mem you can decorate this function to reduce the batch size until no more out-of-memory error occurs.
import torch
from torch_max_mem import maximize_memory_utilization
@maximize_memory_utilization(parameter_name="batch_size")
def knn(x, y, batch_size, k: int = 3):
return torch.cat(
[
torch.cdist(x[start : start + batch_size], y).topk(k=k, dim=0, largest=False).indices
for start in range(0, x.shape[0], batch_size)
],
dim=0,
)
In the code, you can now always pass the largest sensible batch size, e.g.,
x = torch.rand(100, 100, device="cuda")
y = torch.rand(200, 100, device="cuda")
knn(x, y, batch_size=x.shape[0])
🚀
Installation
The most recent release can be installed from PyPI with:
$ pip install torch_max_mem
The most recent code and data can be installed directly from GitHub with:
$ pip install git+https://github.com/mberr/torch-max-mem.git
To install in development mode, use the following:
$ git clone git+https://github.com/mberr/torch-max-mem.git
$ cd torch-max-mem
$ pip install -e .
👐
Contributing
Contributions, whether filing an issue, making a pull request, or forking, are appreciated. See CONTRIBUTING.md for more information on getting involved.
👋
Attribution
Parts of the logic have been developed with Laurent Vermue for PyKEEN.
⚖️
License
The code in this package is licensed under the MIT License.
🍪
Cookiecutter
This package was created with @audreyfeldroy's cookiecutter package using @cthoyt's cookiecutter-snekpack template.
🛠️
For Developers
See developer instrutions
The final section of the README is for if you want to get involved by making a code contribution.
🥼
Testing
After cloning the repository and installing tox with pip install tox, the unit tests in the tests/ folder can be run reproducibly with:
$ tox
Additionally, these tests are automatically re-run with each commit in a GitHub Action.
📖
Building the Documentation
$ tox -e docs
📦
Making a Release
After installing the package in development mode and installing tox with pip install tox, the commands for making a new release are contained within the finish environment in tox.ini. Run the following from the shell:
$ tox -e finish
This script does the following:
- Uses Bump2Version to switch the version number in the
setup.cfgandsrc/torch_max_mem/version.pyto not have the-devsuffix - Packages the code in both a tar archive and a wheel
- Uploads to PyPI using
twine. Be sure to have a.pypircfile configured to avoid the need for manual input at this step - Push to GitHub. You'll need to make a release going with the commit where the version was bumped.
- Bump the version to the next patch. If you made big changes and want to bump the version by minor, you can use
tox -e bumpversion minorafter.
97 Dec 1, 2022
42 Jan 7, 2023
106 Dec 14, 2022
3 Jul 14, 2021
14 Dec 17, 2022
2.5k Jan 6, 2023
529 Jan 3, 2023
6 Jan 12, 2022
4 Jun 20, 2022
12 Dec 19, 2022
3 Jan 21, 2022
119 Jan 07, 2023
36 Dec 02, 2022
215 Dec 25, 2022
3 Dec 29, 2022
76 Jan 03, 2023
4 Nov 28, 2022
5 Nov 16, 2022
516 Dec 28, 2022
262 Dec 21, 2022
3 Jun 15, 2022
34 Nov 07, 2022
6 Nov 18, 2022
13k Jan 09, 2023
6.5k Jan 04, 2023
2.3k Jan 02, 2023
80 Dec 30, 2022
38 Oct 14, 2022
47 Dec 06, 2022