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

Welcome to The Eigensolver Quantum School, a quantum computing crash course designed by students for students.

Code repository accompanying the paper "On Adversarial Robustness: A Neural Architecture Search perspective"

a reccurrent neural netowrk that when trained on a peice of text and fed a starting prompt will write its on 250 character text using LSTM layers

EfficientDet (Scalable and Efficient Object Detection) implementation in Keras and Tensorflow

PyTorch reimplementation of minimal-hand (CVPR2020)

Out-of-Domain Human Mesh Reconstruction via Dynamic Bilevel Online Adaptation

Sample Prior Guided Robust Model Learning to Suppress Noisy Labels

Avatarify Python - Avatars for Zoom, Skype and other video-conferencing apps.

Real-Time and Accurate Full-Body Multi-Person Pose Estimation&Tracking System

Prototype for Baby Action Detection and Classification

Phy-Q: A Benchmark for Physical Reasoning

NeoPlay is the project dedicated to ESport events.

Official Python implementation of the 'Sparse deconvolution'-v0.3.0

The Video-based Accident Detection System built in Python

minimizer-space de Bruijn graphs (mdBG) for whole genome assembly

ISTR: End-to-End Instance Segmentation with Transformers (https://arxiv.org/abs/2105.00637)

VGGFace2-HQ - A high resolution face dataset for face editing purpose

A transformer which can randomly augment VOC format dataset (both image and bbox) online.

Code for "Neural 3D Scene Reconstruction with the Manhattan-world Assumption" CVPR 2022 Oral

Pytorch Implementation of Neural Analysis and Synthesis: Reconstructing Speech from Self-Supervised Representations