江大白 | 3W字长文，带你深入浅出视觉Transformer-下篇(建议收藏!)

由于文章篇幅太长了，这里分成两篇文章“上”和“下”进行学习，“上”篇在https://blog.csdn.net/csdn_xmj/article/details/139482432，这是“下”篇！

本文来源公众号“江大白”，仅用于学术分享，侵权删，干货满满。

原文链接：2W字长文，带你深入浅出视觉Transformer

以下文章来源于知乎：深度眸@知乎

作者：深度眸

链接：https://zhuanlan.zhihu.com/p/308301901

本文仅用于学术分享，如有侵权，请联系后台作删文处理

2 视觉领域的transformer

在理解了标准的transformer后，再来看视觉领域transformer就会非常简单，因为在cv领域应用transformer时候大家都有一个共识：尽量不改动transformer结构，这样才能和NLP领域发展对齐，所以大家理解cv里面的transformer操作是非常简单的。

2.1 分类vision transformer

论文题目：An Image is Worth 16x16 Words:Transformers for Image Recognition at Scale

论文地址：https://arxiv.org/abs/2010.11929

github: https://github.com/lucidrains/vit-pytorch

其做法超级简单，只含有编码器模块：

本文出发点是彻底抛弃CNN，以前的cv领域虽然引入transformer，但是或多或少都用到了cnn或者rnn，本文就比较纯粹了，整个算法几句话就说清楚了，下面直接分析。

2.1.1 图片分块和降维

因为transformer的输入需要序列，所以最简单做法就是把图片切分为patch，然后拉成序列即可。假设输入图片大小是256x256，打算分成64个patch，每个patch是32x32像素

x = rearrange(img, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=p, p2=p)

这个写法是采用了爱因斯坦表达式，具体是采用了einops库实现，内部集成了各种算子，rearrange就是其中一个，非常高效。不懂这种语法的请自行百度。p就是patch大小，假设输入是b,3,256,256，则rearrange操作是先变成(b,3,8x32,8x32)，最后变成(b,8x8,32x32x3)即(b,64,3072)，将每张图片切分成64个小块，每个小块长度是32x32x3=3072，也就是说输入长度为64的图像序列，每个元素采用3072长度进行编码。

考虑到3072有点大，故作者先进行降维：

# 将3072变成dim，假设是1024
self.patch_to_embedding = nn.Linear(patch_dim, dim)
x = self.patch_to_embedding(x)

仔细看论文上图，可以发现假设切成9个块，但是最终到transfomer输入是10个向量，额外追加了一个0和_。为啥要追加？原因是我们现在没有解码器了，而是编码后直接就进行分类预测，那么该解码器就要负责一点点解码器功能，那就是：需要一个类似开启解码标志，非常类似于标准transformer解码器中输入的目标嵌入向量右移一位操作。试下如果没有额外输入，9个块输入9个编码向量输出，那么对于分类任务而言，我应该取哪个输出向量进行后续分类呢？选择任何一个都说不通，所以作者追加了一个可学习嵌入向量输入。那么额外的可学习嵌入向量为啥要设计为可学习，而不是类似nlp中采用固定的token代替？个人不负责任的猜测这应该就是图片领域和nlp领域的差别，nlp里面每个词其实都有具体含义，是离散的，但是图像领域没有这种真正意义上的离散token，有的只是一堆连续特征或者图像像素，如果不设置为可学习，那还真不知道应该设置为啥内容比较合适，全0和全1也说不通。自此现在就是变成10个向量输出，输出也是10个编码向量，然后取第0个编码输出进行分类预测即可。从这个角度看可以认为编码器多了一点点解码器功能。具体做法超级简单，0就是位置编码向量，_是可学习的patch嵌入向量。

# dim=1024
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
# 变成(b,64,1024)
cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)
# 额外追加token，变成b,65,1024
x = torch.cat((cls_tokens, x), dim=1)

2.1.2 位置编码

位置编码也是必不可少的，长度应该是1024，这里做的比较简单，没有采用sincos编码，而是直接设置为可学习，效果差不多

# num_patches=64，dim=1024,+1是因为多了一个cls开启解码标志
self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))

对训练好的pos_embedding进行可视化，如下所示：

相邻位置有相近的位置编码向量，整体呈现2d空间位置排布一样。

将patch嵌入向量和位置编码向量相加即可作为编码器输入

x += self.pos_embedding[:, :(n + 1)]
x = self.dropout(x)

2.1.3 编码器前向过程

作者采用的是没有任何改动的transformer，故没有啥说的。

self.transformer = Transformer(dim, depth, heads, mlp_dim, dropout)

假设输入是(b,65,1024)，那么transformer输出也是(b,65,1024)

2.1.4 分类head

在编码器后接fc分类器head即可

self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, num_classes)
        )

# 65个输出里面只需要第0个输出进行后续分类即可
self.mlp_head(x[:, 0])

到目前为止就全部写完了，是不是非常简单，外层整体流程为：

class ViT(nn.Module):
    def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, channels=3, dropout=0.,emb_dropout=0.):
        super().__init__()
        # image_size输入图片大小 256
        # patch_size 每个patch的大小 32
        num_patches = (image_size // patch_size) ** 2  # 一共有多少个patch 8x8=64
        patch_dim = channels * patch_size ** 2  # 3x32x32=3072
        self.patch_size = patch_size  # 32
        # 1,64+1,1024,+1是因为token，可学习变量，不是固定编码
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        # 图片维度太大了，需要先降维
        self.patch_to_embedding = nn.Linear(patch_dim, dim)
        # 分类输出位置标志，否则分类输出不知道应该取哪个位置
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.dropout = nn.Dropout(emb_dropout)
        # 编码器
        self.transformer = Transformer(dim, depth, heads, mlp_dim, dropout)
        # 输出头
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, num_classes)
        )

    def forward(self, img, mask=None):
        p = self.patch_size

        # 先把图片变成64个patch,输出shape=b,64,3072
        x = rearrange(img, 'b c (h p1) (w p2) -&