Attention for PyTorch with Linear Memory Footprint

Last update: Jan 09, 2022

Overview

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

DLL: Direct Lidar Localization

Pairwise learning neural link prediction for ogb link prediction

Laplacian Score-regularized Concrete Autoencoders

High frequency AI based algorithmic trading module.

Image super-resolution through deep learning

The Pytorch implementation for "Video-Text Pre-training with Learned Regions"

Jremesh-tools - Blender addon for quad remeshing

Improving Query Representations for DenseRetrieval with Pseudo Relevance Feedback:A Reproducibility Study.

Semi-supervised Video Deraining with Dynamical Rain Generator (CVPR, 2021, Pytorch)

Taking A Closer Look at Domain Shift: Category-level Adversaries for Semantics Consistent Domain Adaptation

Keras community contributions

Quickly comparing your image classification models with the state-of-the-art models (such as DenseNet, ResNet, ...)

A Survey on Deep Learning Technique for Video Segmentation

(Py)TOD: Tensor-based Outlier Detection, A General GPU-Accelerated Framework

2021搜狐校园文本匹配算法大赛分比我们低的都是帅哥队

Extreme Rotation Estimation using Dense Correlation Volumes

PointRCNN: 3D Object Proposal Generation and Detection from Point Cloud, CVPR 2019.

😊 Python module for face feature changing

Liquid Warping GAN with Attention: A Unified Framework for Human Image Synthesis

Code for KDD'20 "An Efficient Neighborhood-based Interaction Model for Recommendation on Heterogeneous Graph"

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

DLL: Direct Lidar Localization

Pairwise learning neural link prediction for ogb link prediction

Laplacian Score-regularized Concrete Autoencoders

High frequency AI based algorithmic trading module.

Image super-resolution through deep learning

The Pytorch implementation for "Video-Text Pre-training with Learned Regions"

Jremesh-tools - Blender addon for quad remeshing

Improving Query Representations for DenseRetrieval with Pseudo Relevance Feedback:A Reproducibility Study.

Semi-supervised Video Deraining with Dynamical Rain Generator (CVPR, 2021, Pytorch)

Taking A Closer Look at Domain Shift: Category-level Adversaries for Semantics Consistent Domain Adaptation

Keras community contributions

Quickly comparing your image classification models with the state-of-the-art models (such as DenseNet, ResNet, ...)

A Survey on Deep Learning Technique for Video Segmentation

(Py)TOD: Tensor-based Outlier Detection, A General GPU-Accelerated Framework

2021搜狐校园文本匹配算法大赛 分比我们低的都是帅哥队

Extreme Rotation Estimation using Dense Correlation Volumes

PointRCNN: 3D Object Proposal Generation and Detection from Point Cloud, CVPR 2019.

😊 Python module for face feature changing

Liquid Warping GAN with Attention: A Unified Framework for Human Image Synthesis

Code for KDD'20 "An Efficient Neighborhood-based Interaction Model for Recommendation on Heterogeneous Graph"

2021搜狐校园文本匹配算法大赛分比我们低的都是帅哥队