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.
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()
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,
)
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])
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 .
Contributions, whether filing an issue, making a pull request, or forking, are appreciated. See CONTRIBUTING.md for more information on getting involved.
Parts of the logic have been developed with Laurent Vermue for PyKEEN.
The code in this package is licensed under the MIT License.
This package was created with @audreyfeldroy's cookiecutter package using @cthoyt's cookiecutter-snekpack template.