计算机视觉中自注意力构建块的PyTorch实现
程序员大白
共 5278字,需浏览 11分钟
· 2021-03-07
点击上方“程序员大白”,选择“星标”公众号
重磅干货,第一时间送达
作者:AI Summer
编译:ronghuaiyang
一个非常好用的git仓库,封装了非常全面的计算机视觉中的自注意力构建块,直接调用,无需重复造轮子了。
git仓库地址:https://github.com/The-AI-Summer/self-attention-cv
用einsum和einops在PyTorch中实现计算机视觉的自我注意机制。专注于计算机视觉自注意模块。
使用 pip 安装
$ pip install self-attention-cv
如果你没有GPU,最好是在环境中预装好pytorch。
相关的文章
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 How Positional Embeddings work in Self-Attention
示例代码
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使用ResNet50主干做图像分类
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]
使用Vision Transformer编码器的Unet的复现
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)
位置嵌入可用
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)
参考文献
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.
英文原文:https://github.com/The-AI-Summer/self-attention-cv
推荐阅读
关于程序员大白
程序员大白是一群哈工大,东北大学,西湖大学和上海交通大学的硕士博士运营维护的号,大家乐于分享高质量文章,喜欢总结知识,欢迎关注[程序员大白],大家一起学习进步!
评论
如何画出漂亮的神经网络图?神经网络可视化工具集锦搜集
点击上方“小白学视觉”,选择加"星标"或“置顶”重磅干货,第一时间送达1. draw_convnet一个用于画卷积神经网络的Python脚本https://github.com/gwding/draw_convnet2. NNSVGhttp://alexlenail.me/NN-SVG/LeNet.
小白学视觉
0
如何看待稚晖君的时间管理水平?
点击上方“小白学视觉”,选择加"星标"或“置顶”重磅干货,第一时间送达计算机视觉life整理原文链接:https://www.zhihu.com/question/491456524/answer/2183081310稚晖君究竟是如何安排业余时间去做这么多高水平的项目?而且每个项目的用时也很少,普通
小白学视觉
2
三个优秀的PyTorch实现语义分割框架
点击上方“小白学视觉”,选择加"星标"或“置顶”重磅干货,第一时间送达转自 | 机器学习AI算法工程使用的VOC数据集链接开放在文章中,预训练模型已上传Github,环境我使用Colab pro,大家下载模型做预测即可。代码链接: https://github.com/lixiang007
小白学视觉
0
如何确定神经网络的层数和隐藏层神经元数量?
点击上方“小白学视觉”,选择加"星标"或“置顶”重磅干货,第一时间送达公众号:尤而小屋作者:呦呦鹿鸣编辑:Peter大家好,我是Peter~关于神经网络中隐藏层的层数和神经元个数充满了疑惑。刚好看到一篇文章很好地答疑了,分享给大家~https://zhuanlan.zhihu.com/p/10041
小白学视觉
0
神经网络之CNN与RNN的关系
点击上方“小白学视觉”,选择加"星标"或“置顶”重磅干货,第一时间送达1、CNN介绍CNN是一种利用卷积计算的神经网络。它可以通过卷积计算将原像素很大的图片保留主要特征变成很小的像素图片。本文以李宏毅老师ppt内容展开具体介绍。1.1 Why CNN for Image①为什么引入CNN ?图片示意
小白学视觉
0
“海优最初弄出来可并不是为了给大佬弟子开后门的, 学术界现在几乎不存在任何套利空间,不要相信有什么办法可以投机取巧”
点击上方“小白学视觉”,选择加"星标"或“置顶”重磅干货,第一时间送达知乎问题:为什么部分同学倾向土博+国外博后而不是直接国外读博?感觉身边土博越来越多,基本都是土博➕海博后,申国外phd的明显减少(相比较疫情之前几年)知乎Tianteman: 这个问题下,可以看得出一部分国内硕博信息检索能力差的离
小白学视觉
0
真高!比亚迪员工爆料比亚迪在越南的薪资水平:基本工资480万,全勤奖35万,交通补助20万,餐补110万,每周6天,每天10小时
上一篇:某大公司为逼迫员工离职,竟然把他的工位安排到厕所旁,没想到他直接开始记录领导的如厕时间,还发到公司大群...对此,你怎么看?--完--PS:欢迎在留言区留下你的观点,一起讨论提高。如果今天的文章让你有新的启发,欢迎转发分享给更多人。全文完,感谢你的耐心阅读。如果你还想看到我的文章,请一定给本
开发者全社区
0
太敢穿了!透视纱裙!性感火辣的身材
绝了呀今天的厂花:吴宣仪1995年1月26日,吴宣仪出生于海南省海口市,中国内地流行乐女歌手、影视演员。2016年2月,吴宣仪随宇宙少女发行首张迷你专辑正式出道。2018年4月,她参加《创造101》综艺选秀,获得第二名,成功加入火箭少女101组合。吴宣仪的颜值一直备受称赞,她的五官立体精致,皮肤白皙
逆锋起笔
0