计算机视觉中自注意力构建块的PyTorch实现

首发：AI公园公众号
作者：AI Summer
编译：ronghuaiyang

导读

一个非常好用的git仓库，封装了非常全面的计算机视觉中的自注意力构建块，直接调用，无需重复造轮子了。

git仓库地址：https://github.com/The-AI-Sum...

=================================================================

用einsum和einops在PyTorch中实现计算机视觉的自我注意机制。专注于计算机视觉自注意模块。

使用 pip 安装

$ pip install self-attention-cv

如果你没有GPU，最好是在环境中预装好pytorch。

示例代码

Multi-head attention

import torchfrom self_attention_cv import MultiHeadSelfAttentionmodel = MultiHeadSelfAttention(dim=64)x = torch.rand(16, 10, 64)  # [batch, tokens, dim]mask = torch.zeros(10, 10)  # tokens X tokensmask[5:8, 5:8] = 1y = model(x, mask)

Axial attention

import torchfrom self_attention_cv import AxialAttentionBlockmodel = 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 torchfrom self_attention_cv import TransformerEncodermodel = TransformerEncoder(dim=64,blocks=6,heads=8)x = torch.rand(16, 10, 64)  # [batch, tokens, dim]mask = torch.zeros(10, 10)  # tokens X tokensmask[5:8, 5:8] = 1y = model(x,mask)

Vision Transformer使用ResNet50主干做图像分类

import torchfrom self_attention_cv import ViT, ResNet50ViTmodel1 = ResNet50ViT(img_dim=128, pretrained_resnet=False,                         blocks=6, num_classes=10,                         dim_linear_block=256, dim=256)# ormodel2 = 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]

使用Vision Transformer编码器的Unet的复现

import torchfrom self_attention_cv.transunet import TransUneta = 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 torchfrom self_attention_cv.bottleneck_transformer import BottleneckBlockinp = 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)

位置嵌入可用

1D Positional Embeddings

import torchfrom self_attention_cv.pos_embeddings import AbsPosEmb1D,RelPosEmb1Dmodel = AbsPosEmb1D(tokens=20, dim_head=64)# batch heads tokens dim_headq = 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 torchfrom self_attention_cv.pos_embeddings import RelPosEmb2Ddim = 32  # spatial dim of the feat mapmodel = RelPosEmb2D(    feat_map_size=(dim, dim),    dim_head=128)q = torch.rand(2, 4, dim*dim, 128)y = model(q)

参考文献

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.

—END—

英文原文：https://github.com/The-AI-Sum...

推荐阅读