https://github.com/meakbiyik/torchcache
Cache PyTorch module outputs on-the-fly
https://github.com/meakbiyik/torchcache
cache cache-storage caching caching-library computer-vision llm lvm mru natural-language-processing pytorch torch
Last synced: about 1 year ago
JSON representation
Cache PyTorch module outputs on-the-fly
- Host: GitHub
- URL: https://github.com/meakbiyik/torchcache
- Owner: meakbiyik
- License: mit
- Created: 2023-08-26T22:32:22.000Z (almost 3 years ago)
- Default Branch: main
- Last Pushed: 2024-08-25T12:53:48.000Z (almost 2 years ago)
- Last Synced: 2025-04-12T00:54:57.914Z (about 1 year ago)
- Topics: cache, cache-storage, caching, caching-library, computer-vision, llm, lvm, mru, natural-language-processing, pytorch, torch
- Language: Python
- Homepage: https://torchcache.readthedocs.io
- Size: 258 KB
- Stars: 40
- Watchers: 1
- Forks: 0
- Open Issues: 1
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# torchcache
[](https://github.com/meakbiyik/torchcache/actions/workflows/ci.yaml) [](https://codecov.io/gh/meakbiyik/torchcache) [](https://torchcache.readthedocs.io/en/latest/?badge=latest)
Effortlessly cache PyTorch module outputs on-the-fly with `torchcache`.
Particularly useful for caching and serving the outputs of computationally expensive large, pre-trained PyTorch modules, such as vision transformers. Note that gradients will not flow through the cached outputs.
- [Features](#features)
- [Installation](#installation)
- [Basic usage](#basic-usage)
- [Assumptions](#assumptions)
- [Contribution](#contribution)
## Features
- Cache PyTorch module outputs either in-memory or persistently to disk.
- Simple decorator-based interface for easy usage.
- Uses an MRU (most-recently-used) cache to limit memory/disk usage
## Installation
```bash
pip install torchcache
```
## Basic usage
Quickly cache the output of your PyTorch module with a single decorator:
```python
from torchcache import torchcache
@torchcache()
class MyModule(nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(10, 10)
def forward(self, x):
# This output will be cached
return self.linear(x)
input_tensor = torch.ones(10, dtype=torch.float32)
# Output is cached during the first call...
output = model(input_tensor)
# ...and is retrieved from the cache for the next one
output_cached = model(input_tensor)
```
See documentation at [torchcache.readthedocs.io](https://torchcache.readthedocs.io/en/latest/) for more examples.
## Assumptions
To ensure seamless operation, `torchcache` assumes the following:
- Your module is a subclass of `nn.Module`.
- The module's forward method accepts any number of positional arguments with shapes `(B, *)`, where `B` is the batch size and `*` represents any number of dimensions. All tensors should be on the same device and have the same dtype.
- The forward method returns a single tensor of shape `(B, *)`.
## Contribution
1. Ensure you have Python installed.
2. Install [`poetry`](https://python-poetry.org/docs/#installation).
3. Run `poetry install` to set up dependencies.
4. Run `poetry run pre-commit install` to install pre-commit hooks.
5. Create a branch, make your changes, and open a pull request.