Attention for PyTorch with Linear Memory Footprint

Unofficially implements https://arxiv.org/abs/2112.05682 to get Linear Memory Cost on Attention (+ some sidekick speedup on the GPU when compared to reference implementation in JAX)

Usage:

git clone https://github.com/CHARM-Tx/linear_mem_attention_pytorch
cd linear_mem_attention_pytorch
python setup.py install

Usage:

High Level

from linear_mem_attention_torch.fast_attn import Attention

batch, length, features = 2, 2**8, 64
x, ctx = torch.randn(2, batch, length, features)
mask = torch.randn(batch, length) < 1.

attn = Attention(dim=features, heads = 8, dim_head = 64, bias=False)

# self-attn
v_self = attn(x, x, mask, query_chunk_size=1024, key_chunk_size=4096)

# cross-attn
v_cross = attn(x, ctx, mask, query_chunk_size=1024, key_chunk_size=4096)

Low level

from linear_mem_attention_torch import attention

batch, length, heads, features = 2, 2**8, 8, 64
mask = torch.randn(batch, length) < 1.
q, k, v = torch.randn(3, batch, length, heads, features)

v_ = attention(q, k, v, mask, query_chunk_size=1024, key_chunk_size=4096)

Benchmarks

See examples/example_benchamrk.ipynb for more information.

Citations:

@misc{rabe2021selfattention,
      title={Self-attention Does Not Need $O(n^2)$ Memory}, 
      author={Markus N. Rabe and Charles Staats},
      year={2021},
      eprint={2112.05682},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Attention for PyTorch with Linear Memory Footprint

Related tags

Overview

Attention for PyTorch with Linear Memory Footprint

Usage:

Usage:

High Level

Low level

Benchmarks

Citations:

Owner

A project for developing transformer-based models for clinical relation extraction

This repo will contain code to reproduce and build upon understanding transfer learning

An MQA (Studio, originalSampleRate) identifier for lossless flac files written in Python.

source code the paper Fast and Robust Iterative Closet Point.

A PyTorch implementation of the continual learning experiments with deep neural networks

Industrial knn-based anomaly detection for images. Visit streamlit link to check out the demo.

[PNAS2021] The neural architecture of language: Integrative modeling converges on predictive processing

Simple reference implementation of GraphSAGE.

WSDM2022 Challenge - Large scale temporal graph link prediction

Image-to-image translation with conditional adversarial nets

Computer Vision is an elective course of MSAI, SCSE, NTU, Singapore

Pytorch implementation of Feature Pyramid Network (FPN) for Object Detection

Creative Applications of Deep Learning w/ Tensorflow

[CVPR'2020] DeepDeform: Learning Non-rigid RGB-D Reconstruction with Semi-supervised Data

Pytorch implementation of various High Dynamic Range (HDR) Imaging algorithms

Dogs classification with Deep Metric Learning using some popular losses

Code for IntraQ, PyTorch implementation of our paper under review

Code for "Learning Graph Cellular Automata"

Pytorch Lightning Distributed Accelerators using Ray

repro_eval is a collection of measures to evaluate the reproducibility/replicability of system-oriented IR experiments