目录
摘要
一、介绍
二、相关工作
三、方法
3.1 图像块嵌入 (Patch Embeddings)
3.2 可学习的嵌入 (Learnable Embedding)
3.3 位置嵌入 (Position Embeddings)
3.4 Transformer 编码器
3.5 ViT 张量维度变化举例
3.6 归纳偏置与混合架构
3.7 微调及更高分辨率
3.8 超参数
四、实验
- 论文:https://arxiv.org/abs/2010.11929
- 代码:GitHub - google-research/vision_transformer
- timm:https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
虽然 Transformer 架构已成为 NLP 任务的事实标准,但它在 CV 中的应用仍然有限。在视觉上,注意力要么与卷积网络结合使用,要么用于替换卷积网络的某些组件,同时保持其整体结构。我们证明了这种对 CNNs 的依赖是不必要的,直接应用于图像块序列 (sequences of image patches) 的纯 Transformer 可以很好地执行 图像分类 任务。当对大量数据进行预训练并迁移到多个中小型图像识别基准时 (ImageNet、CIFAR-100、VTAB 等),与 SOTA 的 CNN 相比,Vision Transformer (ViT) 可获得更优异的结果,同时仅需更少的训练资源。
基于自注意力的架构,尤其是 Transformer,已成为 NLP 中的首选模型。主要方法是 在大型文本语料库上进行预训练,然后在较小而特定于任务的数据集上进行微调。 由于 Transformers 的计算效率和可扩展性,训练具有超过 100B 个参数的、前所未有的模型成为了可能。随着模型和数据集的增长,仍未表现出饱和的迹象。
然而,在 CV 中,卷积架构仍然占主导地位。受到 NLP 成功的启发,多项工作尝试将类似 CNN 的架构与自注意力相结合,有些工作完全取代了卷积。后一种模型虽然理论上有效,但由于使用了特定的注意力模式,尚未在现代硬件加速器上有效地扩展。因此,在大规模图像识别中,经典的类 ResNet 架构仍是最先进的。
受 NLP 中 Transformer 成功放缩 (scaling) 的启发,我们尝试将标准 Transformer 直接应用于图像,并尽可能减少修改。为此,我们 将图像拆分为块 (patch),并将这些图像块的线性嵌入序列作为 Transformer 的输入。图像块 image patches 的处理方式与 NLP 应用中的标记 tokens (单词 words) 相同。我们以有监督方式训练图像分类模型。
当在没有强正则化的中型数据集(如 ImageNet)上进行训练时,这些模型产生的准确率比同等大小的 ResNet 低几个百分点。 这种看似令人沮丧的结果可能是意料之中的:Transformers 缺乏 CNN 固有的一些归纳偏置 (inductive biases),例如平移等效性和局部性 (translation equivariance and locality),因此在数据量不足的情况下训练时不能很好地泛化。
但是,如果模型在更大的数据集 (14M-300M 图像) 上训练,情况就会发生变化。我们发现 大规模训练胜过归纳偏置。我们的 Vision Transformer (ViT) 在以足够的规模进行预训练并迁移到具有较少数据点的任务时获得了出色结果。当在公共 ImageNet-21k 数据集或内部 JFT-300M 数据集上进行预训练时,ViT 在多个图像识别基准上接近或击败了最先进的技术。特别是,最佳模型在 ImageNet 上的准确率达到 88.55%,在 ImageNet-RealL 上达到 90.72%,在 CIFAR-100 上达到 94.55%,在 19 个任务的 VTAB 上达到 77.63%。
Transformers 是由 Vaswani 等人提出的 机器翻译 方法,并已成为许多 NLP 任务中最先进的方法。基于大型 Transformers 的模型通常在大型语料库上进行预训练,然后根据手头的任务进行微调:BERT 使用 去噪自监督 预训练任务,而 GPT 工作线使用 语言建模 作为其预训练任务。
应用于图像的简单自注意力要求 每个像素关注所有其他像素。由于像素数量的二次方成本,其无法缩放到符合实际的输入尺寸。因此,曾经有研究者尝试过几种近似方法以便于在图像处理中应用 Transformer。Parmar 等人只在每个 query 像素的局部邻域而非全局应用自注意力,这种局部多头点积自注意力块完全可以代替卷积。在另一种工作中,稀疏 Transformer 采用可放缩的全局自注意力,以便适用于图像。衡量注意力的另一种方法是将其应用于大小不同的块中,在极端情况下仅沿单个轴。许多这种特殊的注意力架构在 CV 任务上显示出很好的效果,但是需要在硬件加速器上有效地实现复杂的工程。
与我们最相关的是 Cordonnier 等人的模型,该模型从输入图像中提取 2×2 大小的块,并在顶部应用完全的自注意力。该模型与ViT 非常相似,但我们的工作进一步证明了 大规模的预训练使普通的 Transformers 能够与 SOTA 的 CNNs 竞争 (甚至更优)。此外,Cordonnier 等人使用 2×2 像素的小块,使模型只适用于小分辨率图像,而我们也能处理中分辨率图像。
将 CNN 与自注意力的形式相结合有很多有趣点,例如增强用于图像分类的特征图,或使用自注意力进一步处理CNN 的输出,如用于目标检测、视频处理、图像分类,无监督目标发现,或统一文本视觉任务。
另一个最近的相关模型是图像 GPT (iGPT),它在降低图像分辨率和颜色空间后对图像像素应用 Transformers。该模型以无监督的方式作为生成模型进行训练,然后可以对结果表示进行微调或线性探测以提高分类性能,在 ImageNet 上达到 72% 的最大精度。
我们的工作增加了在比标准 ImageNet 数据集更大尺度上探索图像识别的论文的数量。使用额外的数据源可以在标准基准上取得 SOTA 的成果。此外,Sun 等人研究了 CNN 性能如何随数据集大小而变化,Kolesnikov、Djolonga 等人从 ImageNet-21k 和JFT-300M 等大规模数据集对 CNN 迁移学习进行了实证研究。我们也关注后两个数据集,但是是训练 Transformers 而非以前工作中使用的基于 ResNet 的模型。
在模型设计中,我们尽可能地遵循原始 Transformer (Vaswani 等, 2017)。 这种有意简单设置的优势在于,可扩展的 NLP Transformer 架构及其高效实现几乎可以开箱即用。
该模型的概述如图 1 所示。标准 Transformer 接受 一维标记嵌入序列 (Sequence of token embeddings) 作为输入。为处理 2D 图像,我们将图像 reshape 为一个展平 (flatten) 的 块序列 ,其中 是原始图像的分辨率, 是通道数 (RGB 图像 ), 是每个图像块的分辨率, 是产生的图像块数,即 Transformer 的有效输入序列长度。Transformer 在其所有层中使用恒定的隐向量 (latent vector) 大小 ,因此我们将图像块展平,并使用 可训练的线性投影 (FC 层) 将维度 映射为 维,同时保持图像块数 不变 (等式 1)。此投影输出称为 图像块嵌入 (Patch Embeddings) (本质就是对每一个展平后的 patch vector 做一个线性变换 / 全连接层 ,由 维降维至 维,得到 ),这好比于 NLP 中的 词嵌入 (Word Embeddings)。图像块嵌入的实现为:
class PatchEmbed(nn.Module):
""" Image to Patch Embedding """
def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
super().__init__()
# (H, W)
img_size = to_2tuple(img_size)
# (P, P)
patch_size = to_2tuple(patch_size)
# N = (H // P) * (W // P)
num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
self.img_size = img_size
self.patch_size = patch_size
self.num_patches = num_patches
# 可训练的线性投影 - 获取输入嵌入
self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
def forward(self, x):
B, C, H, W = x.shape
# FIXME look at relaxing size constraints
assert H == self.img_size[0] and W == self.img_size[1], \
f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*
{self.img_size[1]})."
# (B, C, H, W) -> (B, D, (H//P), (W//P)) -> (B, D, N) -> (B, N, D)
# D=embed_dim=768, N=num_patches=(H//P)*(W//P)
# torch.flatten(input, start_dim=0, end_dim=-1) # 形参:展平的起始维度和结束维度
# 可见 Patch Embedding 操作 3 步到位
x = self.proj(x).flatten(2).transpose(1, 2)
return x
类似于 BERT 的 token,此处 为图像块嵌入序列预设一个 可学习的嵌入,该嵌入在 Transformer 编码器输出的状态/特征 用作 图像表示 (等式 4)。无论是预训练还是微调,都有一个 分类头 (Classification Head) 附加在之后,从而用于图像分类。分类头在预训练时由一个 单层 MLP 实现,在微调时由 单个线性层 实现 (多层感知机与线性模型类似,区别在于 MLP 相对于 FC 层数增加且引入了非线性激活函数,例如 FC + GELU + FC 形式的 MLP)。
更明确地,等式 1 中给长度为的嵌入向量后追加了一个分类向量,用于训练 Transformer 时学习类别信息。假设将图像分为个图像块,输入到 Transformer 编码器中就有个向量,但该取哪一个向量用于分类预测呢?都不合适!一个合理的做法是手动添加一个 可学习的嵌入向量作为用于分类的类别向量 ,同时与其他图像块嵌入向量一起输入到 Transformer 编码器中,最后取追加的首个可学习的嵌入向量作为类别预测结果。所以,追加的首个类别向量可理解为其他个图像块寻找的类别信息。从而,最终输入 Transformer 的嵌入向量总长度为。可学习嵌入 在训练时随机初始化,然后通过训练得到,其具体实现为:
### 随机初始化
self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) # shape = (1, 1, D)
### 分类头 (Classifier head)
self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
### 前馈过程 (Forward)
B = x.shape[0] # Batch Size
# 通过 可学习的线性投影 获取 Input Imgaes 的 Patch Embeddings (实现在 3.1 节)
x = self.patch_embed(x) # x.shape = (B, N, D)
# 可学习嵌入 - 用于分类
cls_tokens = self.cls_token.expand(B, -1, -1) # shape = (B, 1, D)
# 按元素相加 附带 Position Embeddings
x = x + self.pos_embed # shape = (B, N, D) - Python 广播机制
# 按通道拼接 获取 N+1 维 Embeddings
x = torch.cat((cls_tokens, x), dim=1) # shape = (B, N+1, D)
位置嵌入 也被加入图像块嵌入,以保留输入图像块之间的空间位置信息。不同于 CNN,Transformer 需要位置嵌入来编码 patch tokens 的位置信息,这主要是由于 自注意力 的 扰动不变性 (Permutation-invariant),即打乱 Sequence 中 tokens 的顺序并不会改变结果。
相反,若不给模型提供图像块的位置信息,那么模型就需要通过图像块的语义来学习拼图,这就额外增加了学习成本。ViT 论文中对比了几种不同的位置编码方案:
- 无位置嵌入
- 1-D 位置嵌入:考虑把 2-D 图像块视为 1-D 序列
- 2-D 位置嵌入:考虑图像块的 2-D 位置 (x, y)
- 相对位置嵌入:考虑图像块的相对位置
最后发现如果 不提供位置编码效果会差,但其它各种类型的编码效果效果都接近,这主要是因为 ViT 的输入是相对较大的图像块而非像素,所以学习位置信息相对容易很多。
Transformer 原文中默认采用 固定位置编码,ViT 则采用 标准可学习/训练的 1-D 位置编码嵌入,因为尚未观察到使用更高级的 2-D-aware 位置嵌入 (附录 D.4) 能够带来显著的性能提升。在输入 Transformer 编码器之前直接 将图像块嵌入和位置嵌入按元素相加:
# 多 +1 是为了加入上述的 class token
# embed_dim 即 patch embed_dim
self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
# patch emded + pos_embed :图像块嵌入 + 位置嵌入
x = x + self.pos_embed
论文中也对学习到的位置编码进行了可视化,发现相近的图像块的位置编码较相似,且同行或列的位置编码也相近:
关于 Transformer 位置嵌入,详见《【机器学习】详解 Transformer_闻韶-CSDN博客_机器学习transformer》
Transformer 编码器 由交替的 多头自注意力层 (MSA, 附录 A) 和 多层感知机块 (MLP, 等式 2, 3) 构成。在每个块前应用 层归一化 (Layer Norm),在每个块后应用 残差连接 (Residual Connection)。
# MHA
class Attention(nn.Module):
def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
super().__init__()
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim ** -0.5
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
# 附带 dropout
self.proj_drop = nn.Dropout(proj_drop)
def forward(self, x):
B, N, C = x.shape
qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)
attn = (q @ k.transpose(-2, -1)) * self.scale
attn = attn.softmax(dim=-1)
attn = self.attn_drop(attn)
x = (attn @ v).transpose(1, 2).reshape(B, N, C)
x = self.proj(x)
x = self.proj_drop(x)
return x
在 Transformer中,MSA 后跟一个 FFN (Feed-forward network),其包含 两个 FC 层,第一个 FC 将特征从维度 变换成 ,第二个 FC 将特征从维度 恢复成 ,中间的非线性激活函数均采用 GeLU (Gaussian Error Linear Unit,高斯误差线性单元) —— 这实质是一个 MLP (多层感知机与线性模型类似,区别在于 MLP 相对于 FC 层数增加且引入了非线性激活函数,例如 FC + GeLU + FC),实现如下:
class Mlp(nn.Module):
def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
super().__init__()
out_features = out_features or in_features
hidden_features = hidden_features or in_features
self.fc1 = nn.Linear(in_features, hidden_features)
self.act = act_layer()
self.fc2 = nn.Linear(hidden_features, out_features)
self.drop = nn.Dropout(drop)
def forward(self, x):
x = self.fc1(x)
x = self.act(x)
x = self.drop(x)
x = self.fc2(x)
x = self.drop(x)
return x
一个 Transformer Encoder Block 就包含一个 MSA 和一个 FFN,二者都有 跳跃连接 和 层归一化 操作构成 MSA Block 和 MLP Block,实现如下:
# Transformer Encoder Block
class Block(nn.Module):
def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
super().__init__()
# 后接于 MHA 的 Layer Norm
self.norm1 = norm_layer(dim)
# MHA
self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias,
qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
# NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
# 后接于 MLP 的 Layer Norm
self.norm2 = norm_layer(dim)
# 隐藏层维度
mlp_hidden_dim = int(dim * mlp_ratio)
# MLP
self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
def forward(self, x):
# MHA + Add & Layer Norm
x = x + self.drop_path(self.attn(self.norm1(x)))
# MLP + Add & Layer Norm
x = x + self.drop_path(self.mlp(self.norm2(x)))
return x
集合了 类别向量、图像块嵌入 和 位置编码 三者到一体的 输入嵌入向量 后,即可馈入Transformer Encoder。ViT 类似于 CNN,不断前向通过 由 Transformer Encoder Blocks 串行堆叠构成的 Transformer Encoder,最后 提取可学习的类别嵌入向量 —— class token 对应的特征用于 图像分类。整体前向计算过程如下:
- 等式 1:由 图像块嵌入 、类别向量 和位置编码 构成 嵌入输入向量
- 等式 2:由 多头注意力机制、层归一化 和 跳跃连接 (Layer Norm & Add) 构成的 MSA Block,可重复 个,其中第 个输出为
- 等式 3:由 前馈网络 (FFN)、层归一化 和 跳跃连接 (Layer Norm & Add) 构成的 MLP Block,可重复 个,其中第 个输出为
- 等式 4:由 层归一化 (Layer Norm) 和 分类头 (MLP or FC) 输出 图像表示
关于 Transformer 编码器的结构,详见【机器学习】详解 Transformer。
其中,初始输入图像 shape = (b = b, c = 3, h = 256, w = 256) 被切分并展平为:通道数 c = 3、尺寸 P = 32、个数 N = (256×256) / (32×32) = 64 的图像块 (Patch),每个图像块均有 P×P×C = 32×32×3 = 3072 个像素。馈入线性投影层 (相当于通道数由 3 降维为 1) 得到个数为 N = 64、大小 (像素数) 为 D = (32×32×1) = 1024 的图像块嵌入,每个图像块嵌入按元素加 (Element-wise Summary) 入位置向量后,尺寸仍为 N×D = 64×1024,再拼接 (Concat) 一个用于预测分类结果的 1×1024 可学习嵌入向量构成大小为 65×1024 嵌入整体 (长度 N+1 = 64+1 = 65),输入编码器经过一系列前向处理后,得到尺寸仍为 N×D = 65×1024 的输出。
当然,事实上,根据 3.1 节的代码实现,Patch Embedding 只需依次经过 Conv + Flatten + Transpose 操作得到。上述图示和描述仅仅是用于更直观而细粒度地展现馈入 Encoder 的整体 Embedding 的形成过程。
归纳偏置 (Inductive bias):注意到,Vision Transformer 的图像特定归纳偏置比 CNN 少得多。在 CNN 中,局部性、二维邻域结构 和 平移等效性 存在于整个模型的每一层中。而在 ViT 中,只有 MLP 层是局部和平移等变的,因为自注意力层都是全局的。二维邻域结构 的使用非常谨慎:在模型开始时通过将图像切分成块,并在微调时调整不同分辨率图像的位置嵌入 (如下所述)。此外,初始化时的位置嵌入不携带有关图像块的 2D 位置的信息,图像块之间的所有空间关系都必须从头开始学习。
混合架构 (Hybrid Architecture):作为原始图像块的替代方案,输入序列可由 CNN 的特征图构成。在这种混合模型中,图像块嵌入投影 (等式 1) 被用在 经 CNN 特征提取的块 而非 原始输入图像块。作为一种特殊情况,块的空间尺寸可以为 ,这意味着输入序列是通过 简单地将特征图的空间维度展平并投影到 Transformer 维度 获得的。然后,如上所述添加了分类输入嵌入和位置嵌入,再将三者组成的整体馈入 Transformer 编码器。简单来说,就是先用 CNN 提取图像特征,然后由 CNN 提取的特征图构成图像块嵌入。由于 CNN 已经将图像降采样了,所以块尺寸可为 。
关于归纳偏置,详见《【机器学习】浅谈 归纳偏置 (Inductive Bias)_闻韶-CSDN博客_归纳偏置》
通常,我们在大型数据集上预训练 ViT,并对 (更小的) 下游任务进行微调。为此,我们移除了预训练的预测头部,换为一个 零值初始化 的 前馈层 / 全连接层,其中 是下游任务的 类别数。
用比预训练时更高的图像分辨率进行微调通常更有益。当提供更高分辨率的图像时,需要保持图像块大小相同,此时会使有效序列长度更长。Vision Transformer 可处理任意序列长度 (取决于内存限制),但 预训练的位置嵌入可能不再有意义。因此,我们根据它们在原始图像中的位置,对预训练的位置嵌入执行 2D 插值。注意,此分辨率调整和图像块提取是将有关图像 2D 结构的归纳偏置手动注入 Vision Transformer 的唯一点。
def resize_pos_embed(posemb, posemb_new):
# Rescale the grid of position embeddings when loading from state_dict. Adapted from
# https://github.com/google-research/vision_transformer/blob/00883dd691c63a6830751563748663526e811cee/vit_jax/checkpoint.py#L224
_logger.info('Resized position embedding: %s to %s', posemb.shape, posemb_new.shape)
ntok_new = posemb_new.shape[1]
# 除去 class token 的 pos_embed
posemb_tok, posemb_grid = posemb[:, :1], posemb[0, 1:]
ntok_new -= 1
gs_old = int(math.sqrt(len(posemb_grid)))
gs_new = int(math.sqrt(ntok_new))
_logger.info('Position embedding grid-size from %s to %s', gs_old, gs_new)
# 把 pos_embed 变换到 2-D 维度再进行插值
posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
posemb_grid = F.interpolate(posemb_grid, size=(gs_new, gs_new), mode='bilinear')
posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new * gs_new, -1)
posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
return posemb
但这种情形一般会造成性能少许损失,可通过微调模型解决。另外,论文 CPVT 通过 Implicit Conditional Position Encoding 来解决该问题(插入 Conv 来隐式编码位置信息,Zero-padding 让 Conv 学习到绝对位置信息)。
ViT 的超参数主要包括以下,它们直接影响模型参数及计算量:
- Layers:Encoder Block 数量
- Hidden Size D:隐藏层特征大小,其在各 Encoder Block 保持一致
- MLP Size:MLP 特征大小,通常设为 4D
- Heads:MSA 中的 heads 数量
- Patch Size:模型输入的 Patch size,ViT 中共有两个设置:14x14 和 16x16,该参数仅影响计算量
类似 BERT,ViT 原文共定义了 3 种不同大小的模型:Base、Large 和 Huge,其对应的模型参数不同,如下所示。如 ViT-L/16 表示采用 Large 结构,输入 Patch size = 16x16。
事实上,timm 还实现了 Tiny, Small, Base, Large 等多种结构。
ViT 并不像 CNN 那样具有 Inductive Bias,若直接在 ImageNet 上训练,同 level 的 ViT 效果不如 ResNet。但若先在较大的数据集上预训练,然后再对特定的较小数据集进行微调,则效果优于 ResNet。比如 ViT 在Google 私有的 300M JFT 数据集上预训练后,在 ImageNet 上的最好的 Top-1 ACC 可达 88.55%,这在当时已和 ImageNet上的 SOTA 相当了 (Noisy Student EfficientNet-L2 效果为 88.5%,Google 最新的 SOTA 是 Meta Pseudo Labels,效果可达 90.2%):
那么 ViT 至少需要多大的数据量才能比肩 CNN 呢?结果如下图所示。可见预训练的数据量须达到 100M 时才能凸显 ViT 的优势。Transformer 的一个特色其 Scalability:当模型和数据量提升时,性能持续提升。在大数据下,ViT 可能会发挥更大的优势。
Transformer、ResNet 与 Hybrid Transformer 三者的性能变化比较:
此外,论文分析了 不同 Layers 的 Mean Attention Distance,其类比于 CNN 的感受野。结果表明:前面层的 “感受野” 虽然差异很大,但总体相比后面层 “感受野” 较小;而模型后半部分 “感受野” 基本覆盖全局,和 CNN 比较类似,说明 ViT 也最后学习到了类似的范式。
当然,ViT 还可根据 Attention Map 来可视化,得知模型具体关注图像的哪个部分,从结果上看比符合实际:
参考资料:
ViT:视觉Transformer backbone网络ViT论文与代码详解
"未来"的经典之作 ViT:transformer is all you need! - 极市社区
GitHub - dk-liang/Awesome-Visual-Transformer: Collect some papers about transformer with vision. Awesome Transformer with Computer Vision (CV)
ViT:视觉Transformer backbone网络ViT论文与代码详解
用 Vision Transformer 进行图像分类
搞懂 Vision Transformer 原理和代码,看这篇技术综述就够了(二)