Self-attention building blocks for computer vision applications in PyTorch
Implementation of self attention mechanisms for computer vision in PyTorch with einsum and einops. Focused on computer vision self-attention modules.
Install it via pip
It would be nice to install pytorch in your enviroment, in case you don't have a GPU.
pip install self-attention-cv
Related articles
- How Attention works in Deep Learning
- How Transformers work in deep learning and NLP
- How the Vision Transformer (ViT) works in 10 minutes: an image is worth 16x16 words
- Understanding einsum for Deep learning: implement a transformer with multi-head self-attention from scratch
More articles are on the way.
Code Examples
Multi-head attention
import torch
from self_attention_cv import MultiHeadSelfAttention
model = MultiHeadSelfAttention(dim=64)
x = torch.rand(16, 10, 64) # [batch, tokens, dim]
mask = torch.zeros(10, 10) # tokens X tokens
mask[5:8, 5:8] = 1
y = model(x, mask)
Axial attention
import torch
from self_attention_cv import AxialAttentionBlock
model = AxialAttentionBlock(in_channels=256, dim=64, heads=8)
x = torch.rand(1, 256, 64, 64) # [batch, tokens, dim, dim]
y = model(x)
Vanilla Transformer Encoder
import torch
from self_attention_cv import TransformerEncoder
model = TransformerEncoder(dim=64,blocks=6,heads=8)
x = torch.rand(16, 10, 64) # [batch, tokens, dim]
mask = torch.zeros(10, 10) # tokens X tokens
mask[5:8, 5:8] = 1
y = model(x,mask)
Vision Transformer with/without ResNet50 backbone for image classification
import torch
from self_attention_cv import ViT, ResNet50ViT
model1 = ResNet50ViT(img_dim=128, pretrained_resnet=False,
blocks=6, num_classes=10,
dim_linear_block=256, dim=256)
# or
model2 = ViT(img_dim=256, in_channels=3, patch_dim=16, num_classes=10,dim=512)
x = torch.rand(2, 3, 256, 256)
y = model2(x) # [2,10]
A re-implementation of Unet with the Vision Transformer encoder
import torch
from self_attention_cv.transunet import TransUnet
a = torch.rand(2, 3, 128, 128)
model = TransUnet(in_channels=3, img_dim=128, vit_blocks=8,
vit_dim_linear_mhsa_block=512, classes=5)
y = model(a) # [2, 5, 128, 128]
Bottleneck Attention block
import torch
from self_attention_cv.bottleneck_transformer import BottleneckBlock
inp = torch.rand(1, 512, 32, 32)
bottleneck_block = BottleneckBlock(in_channels=512, fmap_size=(32, 32), heads=4, out_channels=1024, pooling=True)
y = bottleneck_block(inp)
Position embeddings are also available
1D Positional Embeddings
import torch
from self_attention_cv.pos_embeddings import AbsPosEmb1D,RelPosEmb1D
model = AbsPosEmb1D(tokens=20, dim_head=64)
# batch heads tokens dim_head
q = torch.rand(2, 3, 20, 64)
y1 = model(q)
model = RelPosEmb1D(tokens=20, dim_head=64, heads=3)
q = torch.rand(2, 3, 20, 64)
y2 = model(q)
2D Positional Embeddings
import torch
from self_attention_cv.pos_embeddings import RelPosEmb2D
dim = 32 # spatial dim of the feat map
model = RelPosEmb2D(
feat_map_size=(dim, dim),
dim_head=128)
q = torch.rand(2, 4, dim*dim, 128)
y = model(q)
References
- Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., ... & Polosukhin, I. (2017). Attention is all you need. arXiv preprint arXiv:1706.03762.
- Wang, H., Zhu, Y., Green, B., Adam, H., Yuille, A., & Chen, L. C. (2020, August). Axial-deeplab: Stand-alone axial-attention for panoptic segmentation. In European Conference on Computer Vision (pp. 108-126). Springer, Cham.
- Srinivas, A., Lin, T. Y., Parmar, N., Shlens, J., Abbeel, P., & Vaswani, A. (2021). Bottleneck Transformers for Visual Recognition. arXiv preprint arXiv:2101.11605.
- Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., ... & Houlsby, N. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929.
93 Jan 05, 2023
21 Apr 09, 2022
0 Jul 18, 2022
2 Nov 23, 2022
308 Jan 05, 2023
17 Dec 14, 2021
21 Sep 06, 2022
116 Dec 27, 2022
138 Oct 28, 2022
3 Nov 06, 2021
4 Dec 22, 2022
364 Dec 28, 2022
19 Dec 06, 2022
35 Dec 06, 2022
58 Oct 11, 2022
70 Dec 04, 2022
975 Dec 30, 2022
6.9k Jan 03, 2023
3k Jan 08, 2023
9 Jun 17, 2022