ViT：视觉Transformer backbone网络ViT论文与代码详解

Visual Transformer

Author：louwill

Machine Learning Lab

今天开始Visual Transformer系列的第一篇文章，主题是Vision Transformer。Vision Transformer (ViT) 可以算是整个Visuier任务的backbone网络。

提出ViT模型的这篇文章题名为An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale，发表于2020年10月份，虽然相较于一些Transformer的视觉任务应用模型 (如DETR) 提出要晚了一些，但作为一个纯Transformer结构的视觉分类网络，其工作还是有较大的开创性意义的。

ViT的总体想法是基于纯Transformer结构来做图像分类任务，论文中相关实验证明在大规模数据集上做完预训练后的ViT模型，在迁移到中小规模数据集的分类任务上以后，能够取得比CNN更好的性能。

ViT模型详解

ViT模型整体结构概览如图1所示。

ViT的核心流程包括图像分块处理 (make patches)、图像块嵌入 (patch embedding)与位置编码、Transformer编码器和MLP分类处理等4个主要部分。下面分别从这四个流程部分来阐述ViT的基本设计。

图像分块处理 (make patches)

第一步可以看作是一个图像预处理步骤。在CNN中，直接对图像进行二维卷积处理即可，不需要特殊的预处理流程。但Transformer结构不能直接处理图像，在此之前需要对其进行分块处理。

假设一个图像x∈H×W×C，现在将其分成P×P×C的patches，那么实际有N=HW/P2个patches，全部patches的维度就可以写为N×P×P×C。然后将每个patch进行展平，相应的数据维度就可以写为N×(P2×C)。这里N可以理解为输入到Transformer的序列长度，C为输入图像的通道数，P为图像patch的大小。

图像块嵌入 (patch embedding)

图像分块仅仅是一道预处理流程，要将N×(P2×C)的向量维度，转化为N×D大小的二维输入，还需要做一个图像块嵌入的操作，类似NLP中的词嵌入，块嵌入也是一种将高维向量转化为低维向量的方式。

所谓图像块嵌入，其实就是对每一个展平后的patch向量做一个线性变换，即全连接层，降维后的维度为D。

上式中的E即为块嵌入的全连接层，其输入大小为(P2×C)，输出大小为D。

值得注意的是，上式中给长度为N的向量还追加了一个分类向量，用于Transformer训练过程中的类别信息学习。假设将图像分为9个patch，即N=9，输入到Transformer编码器中就有9个向量，但对于这9个向量而言，该取哪一个向量做分类预测呢？取哪一个都不合适。一个合理的做法就是人为添加一个类别向量，该向量是可学习的嵌入向量，与其他9个patch嵌入向量一起输入到Transformer编码器中，最后取第一个向量作为类别预测结果。所以，这个追加的向量可以理解为其他9个图像patch寻找的类别信息。

位置编码 (position encoding)

为了保持输入图像patch之间的空间位置信息，还需要对图像块嵌入中添加一个位置编码向量，如上式中的Epos所示，ViT的位置编码没有使用更新的2D位置嵌入方法，而是直接用的一维可学习的位置嵌入变量，原先是论文作者发现实际使用时2D并没有展现出比1D更好的效果。

ViT前向流程

集合了类别向量追加、图像块嵌入和位置编码为一体的嵌入输入向量后，就可以直接进入Transformer编码器部分了，主要包括MSA和MLP两个部分。所以，ViT的编码器前向计算过程可以归纳如下：

第一个式子即前述的图像块嵌入、类别向量追加和位置编码；第二个式子为MSA部分，包括多头自注意力、跳跃连接 (Add) 和层规范化 (Norm) 三个部分，可以重复L个MSA block；第三个式子为MLP部分，包括前馈网络 (FFN)、跳跃连接 (Add) 和层规范化 (Norm) 三个部分，也可以重复L个MSA block。第四个式子为层规范化。最后以一个MLP作为分类头 (Classification Head)。

为了更加清晰的展示ViT模型结构和训练过程中的向量变化，下图给出了ViT的向量维度变化图。

图来自于极市平台

ViT训练与实验

ViT训练方法

ViT的基本训练策略是在大数据集上先做预训练，然后在在小数据集上做迁移使用。ViT做预训练使用到的大数据集包括：

• ILSVRC-2012 ImageNet dataset：1000 classes
• ImageNet-21k：21k classes
• JFT：18k High Resolution Images
  

其中JFT是一个谷歌的内部大规模图像数据集，约有300M图像18291个类别标注。

ViT预训练迁移到的数据集包括：

• CIFAR-10/100
• Oxford-IIIT Pets
• Oxford Flowers-102
• VTAB

论文共设计了Base、Large和Huge三款不同大小的ViT模型，分别表示基础模型、大模型和超大模型，三款模型的各参数如下表所示。

比如说，ViT-B/16就表示patch size为16的ViT-Base模型。

ViT实验设计

ViT最核心的实验就是将前述的训练方法进行实现，即在大规模数据集上预训练后迁移到小数据集上看模型效果。为了比对CNN模型，论文特地用了Big Transfer (BiT)，该模型使用大的ResNet进行监督迁移学习，是2020 ECCV上提出的一个大CNN模型。另外一个比对CNN模型是2020年CVPR上的Noisy Student模型，是一个半监督的大型CNN模型。

ViT、BiT和Nosiy Student模型经三大数据集预训练后在各小数据集上的准确率如下表所示。

可以看到，ViT经过大数据集的预训练后，在各小数据集上的迁移后准确率超过了一些SOTA CNN模型的结果。但要取得这种超越CNN的性能效果，需要大的预训练数据集和大模型的结合。

所以第二个实验就是ViT对预训练数据集规模到底有怎样的要求？论文针对此问题做了一个对比实验。分别在ImageNet、ImageNet-21k和JFT-300M进行预训练，三个数据集规模分别为小数据集、中等规模数据集和超大数据集，预训练效果如下图所示。

从图中可以看到，在最小的数据集ImageNet上进行预训练时，尽管作者加了大量的正则化操作，ViT-Large模型性能不如ViT-base模型，更远不如BiT的性能。在中等规模的ImageNet-21k数据集上，大家的表现都差不多，只有到了JFT-30M这样的超大数据集上，ViT模型才能发挥出它的优势和效果。

总而言之，大的预训练数据集加上大模型，是ViT取得SOTA性能的关键因素。

ViT代码使用与解读

ViT模型实现目前已经有开源的框架vit-pytorch可以直接调用，直接pip安装即可：

pip install vit-pytorch

vit-pytorch用法如下：

import torch
from vit_pytorch import ViT
# 创建ViT模型实例
v = ViT(
    image_size = 256,
    patch_size = 32,
    num_classes = 1000,
    dim = 1024,
    depth = 6,
    heads = 16,
    mlp_dim = 2048,
    dropout = 0.1,
    emb_dropout = 0.1
)
# 随机化一个图像输入
img = torch.randn(1, 3, 256, 256)
# 获取输出
preds = v(img) # (1, 1000)

各参数含义分别为：

• image_size：原始图像尺寸

• patch_size：图像块的尺寸

• num_classes：类别数量

• dim：Transformer隐变量维度大小

• depth：Transformer编码器层数

• Heads：MSA中的head数

• dropout：失活比例

• emb_dropout：嵌入层失活比例

下面我们重点看一下vit.py的代码解读。ViT以Attention和Transformer为基础，所以搭建逻辑跟Transformer是一样的，先把底层各组件搭建好后，按照ViT的前向流程进行封装即可。ViT所需的底层搭建组件包括规范化层、FFN、Attention，然后在此三个组件基础上搭建Transformer，最后基于Transformer和ViT前向流程搭建ViT。下面我们分三个步骤来看ViT的搭建过程。

(1) 底层组件规范化层、FFN、Attention

# 导入相关模块
import torch
from torch import nn, einsum
import torch.nn.functional as F
from einops import rearrange, repeat
from einops.layers.torch import Rearrange

# 辅助函数，生成元组
def pair(t):
    return t if isinstance(t, tuple) else (t, t)

# 规范化层的类封装
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)
# FFN
class FeedForward(nn.Module):
    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)
# Attention
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.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):
        b, n, _, h = *x.shape, self.heads
        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 = h), qkv)
        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
        attn = self.attend(dots)
        out = einsum('b h i j, b h j d -> b h i d', attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

(2) 搭建Transformer

# 基于PreNorm、Attention和FFN搭建Transformer
class Transformer(nn.Module):
    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 x

(3) 搭建ViT

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_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.'
        # patch数量
        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)'
        # 定义块嵌入
        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),
        )
        # 定义位置编码
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
        # 定义类别向量
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
        self.dropout = nn.Dropout(emb_dropout)

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

        self.pool = pool
        self.to_latent = nn.Identity()
        # 定义MLP
        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, num_classes)
        )
        
    # ViT前向流程
    def forward(self, img):
        # 块嵌入
        x = self.to_patch_embedding(img)
        b, n, _ = x.shape
        # 追加类别向量
        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
        x = torch.cat((cls_tokens, x), dim=1)
        # 追加位置编码
        x += self.pos_embedding[:, :(n + 1)]
        # dropout
        x = self.dropout(x)
        # 输入到transformer
        x = self.transformer(x)
        x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]
        x = self.to_latent(x)
        # MLP
        return self.mlp_head(x)

小结

ViT作为Visual Transformer的一篇开创性研究，可以算是了解该方向的一篇必读论文了。今年上半年以来，大量基于ViT的视觉任务研究不断的被提出，ViT在其中基本上扮演了类似VGG16或者ResNet-52在CNN中Backbone的角色。虽然是一篇开创性的工作，但ViT仍有大量的使用限制，大数据集和大模型，这两点就已经将大多数人望而却步了。当然，这些缺陷，在后来的研究中也在不断的被克服。

参考资料：

An Image Is Worth 16X16 Words: Transformers for Image Recognition at Scale

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

https://mp.weixin.qq.com/s/ozUHHGMqIC0-FRWoNGhVYQ

往期精彩：

求个在看