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

Attempt at implementation of a simple GAN using Keras

A High-Level Fusion Scheme for Circular Quantities published at the 20th International Conference on Advanced Robotics

Azion the best solution of Edge Computing in the world.

Joint-task Self-supervised Learning for Temporal Correspondence (NeurIPS 2019)

CCPD: a diverse and well-annotated dataset for license plate detection and recognition

A lightweight deep network for fast and accurate optical flow estimation.

Orbivator AI - To Determine which features of data (measurements) are most important for diagnosing breast cancer and find out if breast cancer occurs or not.

Segmentation-Aware Convolutional Networks Using Local Attention Masks

Prediction of MBA refinance Index (Mortgage prepayment)

Unofficial keras(tensorflow) implementation of MAE model from Masked Autoencoders Are Scalable Vision Learners

Abstractive opinion summarization system (SelSum) and the largest dataset of Amazon product summaries (AmaSum). EMNLP 2021 conference paper.

Implementation of CVPR 2020 Dual Super-Resolution Learning for Semantic Segmentation

Generative Query Network (GQN) in PyTorch as described in "Neural Scene Representation and Rendering"

This is the official Pytorch implementation of "Lung Segmentation from Chest X-rays using Variational Data Imputation", Raghavendra Selvan et al. 2020

Implementations of the algorithms in the paper Approximative Algorithms for Multi-Marginal Optimal Transport and Free-Support Wasserstein Barycenters

Quantile Regression DQN a Minimal Working Example, Distributional Reinforcement Learning with Quantile Regression

Understanding Convolution for Semantic Segmentation

PyTorch implementations for our SIGGRAPH 2021 paper: Editable Free-viewpoint Video Using a Layered Neural Representation.

Open source simulator for autonomous vehicles built on Unreal Engine / Unity, from Microsoft AI & Research

Official PyTorch implementation of the paper "Self-Supervised Relational Reasoning for Representation Learning", NeurIPS 2020 Spotlight.