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

Python script that allows you to automatically setup your Growtopia server.

REGTR: End-to-end Point Cloud Correspondences with Transformers

An implementation of the "Attention is all you need" paper without extra bells and whistles, or difficult syntax

Multi-Stage Progressive Image Restoration

FaceQgen: Semi-Supervised Deep Learning for Face Image Quality Assessment

This repo holds the code of TransFuse: Fusing Transformers and CNNs for Medical Image Segmentation

An implementation of chunked, compressed, N-dimensional arrays for Python.

Java and SHACL code commented in the paper "Towards compliance checking in reified I/O logic via SHACL" submitted to ICAIL 2021

Framework for joint representation learning, evaluation through multimodal registration and comparison with image translation based approaches

Implementation of: "Exploring Randomly Wired Neural Networks for Image Recognition"

Cereal box identification in store shelves using computer vision and a single train image per model.

SciKit-Learn Laboratory (SKLL) makes it easy to run machine learning experiments.

Joint project of the duo Hacker Ninjas

Extracting and filtering paraphrases by bridging natural language inference and paraphrasing

A Library for Modelling Probabilistic Hierarchical Graphical Models in PyTorch

Two-stage CenterNet

From this paper "SESNet: A Semantically Enhanced Siamese Network for Remote Sensing Change Detection"

[SIGGRAPH 2022 Journal Track] AvatarCLIP: Zero-Shot Text-Driven Generation and Animation of 3D Avatars

Code for the prototype tool in our paper "CoProtector: Protect Open-Source Code against Unauthorized Training Usage with Data Poisoning".

A basic implementation of Layer-wise Relevance Propagation (LRP) in PyTorch.