《AN IMAGE IS WORTH 16X16 WORDS:
TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE》
论文共有22页，表格和图像很多，网络模型结构解释的很清楚，并且用四个公式展示了模型的计算过程；本文章对其进行精度，并对源码进行剖析，希望读者可以耐心读下去。
论文地址：https://arxiv.org/abs/2010.11929
源码地址（pytorch）：https://github.com/lucidrains/vit-pytorch
目录
虽然Transformer架构已经成为自然语言处理任务的事实上的标准，但它在计算机视觉中的应用仍然有限。在视觉上，注意力要么与卷积网络结合使用，要么用于替换卷积网络的某些组成部分，同时保持其整体结构的到位。我们表明，这种对CNN的依赖是不必要的，直接应用于图像块序列的纯Transformer可以很好地执行图像分类任务。当对大量数据进行预训练并转移到多个中小规模的图像识别基准时(ImageNet、CIFAR-100、VTAB等)，Vision Transformer ( ViT ) 相对于先进的卷积网络获得了优异的结果，同时训练所需的计算资源也大大减少。
我们探索了Transformer在图像识别中的直接应用。不同于以往在计算机视觉中使用自注意力的工作，我们除了初始的patches提取步骤外，并没有引入图像特定的诱导偏差。相反，我们将图像理解为一个patches序列，并通过标准的 Transformer 编码器对其进行处理，如同在NLP中使用一样。这种简单而又可扩展的策略在结合大规模数据集的预训练时表现出惊人的效果。因此，Vision Transformer 在许多图像分类数据集上匹配或超过了艺术状态，同时相对于训练前来说相对便宜。
虽然这些初步成果令人鼓舞，但仍存在许多挑战。一是将ViT应用到其他计算机视觉任务中，如检测和分割。我们的结果，再加上 《End-to-end object detection with transformers》(2020) 的结果，表明这一方法的前景。另一个挑战是继续探索自监督的预训练方法。我们的初始实验从自监督预训练方面表现出了改进，但自监督与大规模监督预训练相比仍有较大差距。最后，ViT 的进一步缩放很可能导致性能的提高。
原文：我们将图像分割成固定大小的块，线性嵌入其中的每个块，加入位置嵌入，并将得到的向量序列反馈给标准的Transformer编码器。为了执行分类，我们采用在序列中添加额外可学习的“classification token”的标准方法。
结合原文，把模型的结构解析成以下三个步骤：
（1）把大小为256*256的图像切割成16*16的patches，每一个patch的像素大小是14*14
（2）对一张图像进行切割、位置嵌入、分类嵌入的操作，返回一个大小为197*768的矩阵，准备送入到标准的Transformer编码器中。
（3）本文章举例处理了一张图像，即batch_size=1。[1, 197, 768]经标准的Transformer编码器输出[1, 197, 768]，然后经过分类头，输出分类结果。
以上三个步骤的计算公式，可以由下面的四个公式概括：
如下图所示，右侧是ViT模型的结构图，左边是对各个公式的参数解释，其对应右侧的各个模型块
构建ViT模型的结构前，需要搭建一个标准的Transformer编码器结构。
一个标准的Transformer编码器结构包括：多头注意力机制、层规范化和残差连接和多层感知机（简单的神经网络）。
******如果对Transformer的结构不是很熟悉，可以参考我的 “变形金刚 Transformer” 专栏******
层规范化：
class PreNorm(nn.Module):
"""层规范化"""
def __init__(self, dim, fn):
super().__init__()
self.norm = nn.LayerNorm(dim)
self.fn = fn
def forward(self, x, **kwargs):
return self.fn(self.norm(x), **kwargs)多层感知机：
class FeedForward(nn.Module):
"""前馈神经网络（MLP）"""
def __init__(self, dim, hidden_dim, dropout = 0.):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(hidden_dim, dim),
nn.Dropout(dropout)
)
def forward(self, x):
return self.net(x)多头注意力机制：
class Attention(nn.Module):
"""多头注意力机制"""
def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
super().__init__()
inner_dim = dim_head * heads
project_out = not (heads == 1 and dim_head == dim)
self.heads = heads
self.scale = dim_head ** -0.5
self.attend = nn.Softmax(dim = -1)
self.dropout = nn.Dropout(dropout)
self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
self.to_out = nn.Sequential(
nn.Linear(inner_dim, dim),
nn.Dropout(dropout)
) if project_out else nn.Identity()
def forward(self, x):
qkv = self.to_qkv(x).chunk(3, dim = -1)
q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
attn = self.attend(dots)
attn = self.dropout(attn)
out = torch.matmul(attn, v)
out = rearrange(out, 'b h n d -> b n (h d)')
return self.to_out(out)Transformer编码器：
class Transformer(nn.Module):
"""标准的Transformer编码器"""
def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
super().__init__()
self.layers = nn.ModuleList([])
for _ in range(depth):
self.layers.append(nn.ModuleList([
PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout)),
PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
]))
def forward(self, x):
# 实现残差网络的相加
for attn, ff in self.layers:
x = attn(x) + x
x = ff(x) + x
return xViT模型结构可以分成三个部分：
class ViT(nn.Module):
def __init__(self, *,
image_size,
patch_size,
num_classes,
dim,
depth,
heads,
mlp_dim,
pool='cls',
channels=3,
dim_head=64,
dropout=0.,
emb_dropout=0.
):
super().__init__()
# 图像尺寸
image_height, image_width = pair(image_size)
# patch尺寸
patch_height, patch_width = pair(patch_size)
assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'
# patches数量
num_patches = (image_height // patch_height) * (image_width // patch_width)
# 一个patch的维度
patch_dim = channels * patch_height * patch_width
assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'
# 处理原始图像
# Rearrange：按照字符串的含义对目标进行重新排列的操作
# img = torch.randn(1, 3, 256, 256) '1 3 (256 32) (256 32) -> 1 (256 256) (32 32 3)'
# 分成32*32个patch，并平铺成1024个patch，每一个patch的大小为8*8
# 然后经过一个全连接层(3072, 1024)输出处理后的图像(三维：32*32*3=3072)
self.to_patch_embedding = nn.Sequential(
Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width),
nn.Linear(patch_dim, dim),
)
# Sequential(
# (0): Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=32, p2=32)
# (1): Linear(in_features=3072, out_features=1024, bias=True)
# )
# print(self.to_patch_embedding)
# 位置嵌入（patches+cls的位置信息）
self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
# class token
self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
self.dropout = nn.Dropout(emb_dropout)
# 定义标准Transformer模型结构
self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)
self.pool = pool
# 不区分参数的占位符标识运算符。
# identity模块不改变输入，直接return input
# 一种编码技巧吧，比如我们要加深网络，有些层是不改变输入数据的维度的，这时就可以使用此函数
# 这个网络层的设计是仅用于占位的，即不干活，只是有这么一个层，放到残差网络里就是在跳过连接的地方用这个层，显得没有那么空虚！
self.to_latent = nn.Identity()
# 定义MLP分类头的模型结构
# 首先经过LN，然后经过一个全连接层(1024, 1000)，输出分类结果
self.mlp_head = nn.Sequential(
nn.LayerNorm(dim),
nn.Linear(dim, num_classes)
)
def forward(self, img):
# img = torch.randn(1, 3, 256, 256)
# 即往网络中送入一张256*256的三维图像(batch_size=1)
x = self.to_patch_embedding(img)
# 平铺后，共有8*8个patch，其大小为32*32
print('平铺:', x.size()) # torch.Size([1, 64, 1024])
b, n, _ = x.shape
# print(b, n) # 1 64
# class tokens
cls_tokens = repeat(self.cls_token, '1 n d -> b n d', b=b)
# print(cls_tokens.size()) # torch.Size([1, 1, 1024])
x = torch.cat((cls_tokens, x), dim=1) # 沿着dim=1方向对cls_token和x进行拼接
x += self.pos_embedding[:, :(n + 1)] # 拼接后的x与嵌入未知信息
x = self.dropout(x)
print(x.size()) # torch.Size([1, 65, 1024])
# 把预处理的图像送入Transformer模型中
x = self.transformer(x)
print(x.size()) # torch.Size([1, 65, 1024])
# 如果pool=='mean'，返回dim=1方向上的元素平均值
# 否则，直接返回dim=0方向上的第一行的所有元素，即class tokens
x = x.mean(dim=1) if self.pool == 'mean' else x[:, 0]
print(x.size()) # torch.Size([1, 1024])
# 不区分参数的占位符标识运算符。
x = self.to_latent(x)
# 返回分类头的分类结果
return self.mlp_head(x)测试代码：输入指定参数，调用ViT模型，输出分类矩阵：
输入参数含义：
- image_size=256, # 图像尺寸
- patch_size=32, # patches大小
- num_classes=1000, # 标签数量
- dim=1024, # patch的维度
- depth=6, # 模型深度（编码器的数量）
- heads=16, # 注意力头的数量
- mlp_dim=2048, # Transformer中MLP的输出维度
- dropout=0.1, # Transformer中MLP的舍弃率
- emb_dropout=0.1 # 嵌入舍弃率
import torch
from vit_pytorch import ViT
def test():
v = ViT(
image_size=256, # 图像尺寸
patch_size=32, # patches数量
num_classes=1000,# 标签数量
dim=1024, # patch的维度
depth=6, # 模型深度（编码器的数量）
heads=16, # 注意力头的数量
mlp_dim=2048, # Transformer中MLP的输出维度
dropout=0.1, # Transformer中MLP的舍弃率
emb_dropout=0.1 # 嵌入舍弃率
)
img = torch.randn(1, 3, 256, 256)
preds = v(img)
# 如果preds.shape ！= (1, 1000)，自行中断程序，并报错：'correct logits outputted'
assert preds.shape == (1, 1000), 'correct logits outputted'
return preds.shape
VIT_result = test()
print(VIT_result) # torch.Size([1, 1000])
下面的表格表示的是ViT的模型变体，分别是ViT-B、ViT-L、ViT-H。
我们基于用于BERT《BERT: Pre-training of deepbidirectional transformers for language understanding》的ViT配置，如下表所示。“Base”和“Large”模型直接采用BERT，我们加入了较大的“Huge”模型。
下面我们使用简要说明来表示模型大小和输入patch大小：例如，ViT-L / 16表示输入patch大小为16×16的 “Large” 变体。注意，Transformer的序列长度与patch大小的平方成反比，因此具有较小patch大小的模型计算代价更高。
我们发现，大规模的训练会克服归纳性偏差。我们的视觉变形金刚( ViT )在足够规模的预训练后转移到数据量较少的任务时取得了很好的效果。在公共ImageNet-21k数据集或内部 JFT-300M 数据集上进行预训练时，ViT接近器在多个图像识别基准上击败了艺术状态。特别是最佳模型在ImageNet上达到了88.55%的准确率，在ImageNet-ReaL上达到了90.72%，在CIFAR-100上达到了94.55%。
在JFT - 300M上预训练的较小的 ViT - L / 16 模型在所有任务上都优于 BiT - L ( 在同一数据集上进行预训练 )模型，同时训练所需的计算资源大大减少。更大的模型 ViT-H / 14 进一步提高了性能，特别是在更具挑战性的数据集- ImageNet、CIFAR-100和VTAB套件上。有趣的是，该模型对预训练的计算量仍然大大低于现有的技术状态。但是，我们注意到，预训练效率不仅可能受到体系结构选择的影响，还可能受到其他参数的影响，如训练计划、优化器、权值衰减等。最后，在公开的ImageNet-21k数据集上预训练的 ViT-L / 16 模型在大多数数据集上也表现良好，同时占用了很少的资源进行预训练：它需要使用8个核心的标准云TPUv3训练大约30天。
下面的表格表示微调不同ResNet中Adam、SGD的消融实验：
预训练的模型集：
如下图所示：比较不同架构的性能与预训练计算：ViT，ResNets和混合Transformer的比较。在相同的计算开销下，ViT的性能一般优于ResNet。对于较小的模型尺寸，混合Transformer比纯Transformer有所改善，而对于较大的模型尺寸，纯Transformer比混合Transformer有所改善。
启发：这个结果有些令人惊讶，因为人们可能期望卷积局部特征处理能够在任何大小上辅助ViT。其次，ViT似乎没有在所尝试的范围内饱和，从而激励了未来的缩放努力（主要的缩放策略便是对模型的宽度（w）、深度（d）和分辨率（r）进行调整）。
为了了解ViT如何使用自注意来整合图像中的信息，我们分析了不同层次的注意权重所跨越的平均距离。这种 “注意力距离” 类似于CNN的感受野大小。平均注意距离在较低层的头之间是高度可变的，有的头会关注图像的大部分，有的头会关注查询位置处或附近的小区域。随着深度的增加，所有注意力头的注意距离增加。在网络的后半部分，大多数的heads通过tokens广泛参加。
如下图所示，按heads和网络深度划分参加区域的规模。通过平均查询像素与所有其他像素之间的距离，加权得到的注意力权重，对128幅样本图像计算了注意距离。每一个点显示了16个heads中的1个head在1层Transformer中的平均注意距离。图像宽度为224像素。
分类头：这个设计继承了Transformer模型的文本，我们在整个主论文中使用它。最初尝试只使用图像块嵌入，全局平均池( GAP )它们，随后使用线性分类器（就像ResNet的最终特征映射一样）的性能很差。但是，我们发现性能差既不是由于额外的token，也不是GAP操作造成的。相反，他在表现上的差异完全由不同的学习率影响。
下表总结了本消融研究在 ViT-B / 16 模型上的结果。正如我们所看到的，虽然没有位置嵌入的模型和有位置嵌入的模型的性能有很大的差距，但是不同的位置信息编码方式之间几乎没有差别。我们推测，由于我们的Transformer编码器工作在 patch-level 的输入上，相对于像素级，在如何编码空间信息方面的差异并不重要。更确切地说，在 patch-level 输入中，空间维度远小于原始的像素级输入，如14 × 14而不是224 × 224，对于这些不同的位置编码策略，学习表示空间关系同样容易。
即便如此，网络学习到的位置嵌入相似度的具体模式取决于训练超参数。
认真阅读完毕，定会有所收获！
>>>如有疑问，欢迎评论区一起探讨！