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
9 Dec 23, 2022
293 Dec 28, 2022
76 Dec 13, 2022
180 Dec 22, 2022
50 Jun 07, 2022
4 Aug 19, 2022
122 Dec 13, 2022
95 Nov 25, 2022
36 Sep 07, 2022
4 Nov 23, 2022
8 Dec 24, 2022
1.6k Dec 30, 2022
1.1k Dec 30, 2022
46 Dec 03, 2022
1 Jan 06, 2022
63 Nov 14, 2022
14 Dec 08, 2022
68 Jan 02, 2023
38 Sep 16, 2022
13 Dec 28, 2022