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
1.3k Jan 03, 2023
0 Oct 13, 2021
1 Dec 12, 2021
76 Dec 13, 2022
3 Nov 24, 2022
574 Jan 02, 2023
89 Dec 07, 2022
118 Dec 26, 2022
13 Oct 19, 2022
92 Dec 28, 2022
14 Dec 17, 2021
77 Nov 22, 2022
656 Dec 28, 2022
242 Jan 07, 2023
4.1k Dec 28, 2022
3.1k Dec 28, 2022
39 Nov 23, 2022
145 Sep 16, 2022
51 Dec 10, 2022
1.4k Jan 01, 2023