从这篇文章开始，我们把深度学习语义分割的目光转向Transformer，即基于ViT的语义分割模型。在正式介绍Transformer分割网络之前，需要先了解一下ViT的分类网络。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-pytorchvit-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仍有大量的使用限制，大数据集和大模型，这两点就已经将大多数人望而却步了。当然，这些缺陷，在后来的研究中也在不断的被克服。
往期精彩：
深度学习论文精读[13]：Deeplab v3+
深度学习论文精读[12]：Deeplab v3
深度学习论文精读[11]：Deeplab v2
深度学习论文精读[10]：Deeplab v1
深度学习论文精读[9]：PSPNet
深度学习论文精读[8]：ParseNet
深度学习论文精读[7]：nnUNet
深度学习论文精读[6]：UNet++
深度学习论文精读[5]：Attention UNet
深度学习论文精读[4]：RefineNet
深度学习论文精读[3]：SegNet
深度学习论文精读[2]：UNet网络
深度学习论文精读[1]：FCN全卷积网络