视觉Transformer经典论文——ViT、DeiT的与原理解读与实现

视觉Transformer经典论文——ViT、DeiT的与原理解读与实现

最近ChatGPT、文心一言等大模型爆火，追究其原理还是绕不开2017年提出的Transformer结构。Transformer算法自从提出后，在各个领域的相关工作还是非常多的，这里分享之前在其他平台的一篇笔记给大家，详细解读CV领域的两个经典Transformer系列工作——ViT和DeiT。

ViT算法综述

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

之前的算法大都是保持CNN整体结构不变，在CNN中增加attention模块或者使用attention模块替换CNN中的某些部分。ViT算法中，作者提出没有必要总是依赖于CNN，仅仅使用Transformer结构也能够在图像分类任务中表现很好。

受到NLP领域中Transformer成功应用的启发，ViT算法中尝试将标准的Transformer结构直接应用于图像，并对整个图像分类流程进行最少的修改。具体来讲，ViT算法中，会将整幅图像拆分成小图像块，然后把这些小图像块的线性嵌入序列作为Transformer的输入送入网络，然后使用监督学习的方式进行图像分类的训练。ViT算法的整体结构如 图1 所示。

图1：ViT算法结构示意图

该算法在中等规模（例如ImageNet）以及大规模（例如ImageNet-21K、JFT-300M）数据集上进行了实验验证，发现：

• Tranformer相较于CNN结构，缺少一定的平移不变性和局部感知性，因此在数据量不充分时，很难达到同等的效果。具体表现为使用中等规模的ImageNet训练的Tranformer会比ResNet在精度上低几个百分点。
• 当有大量的训练样本时，结果则会发生改变。使用大规模数据集进行预训练后，再使用迁移学习的方式应用到其他数据集上，可以达到或超越当前的SOTA水平。

图2 为大家展示了使用大规模数据集预训练后的 ViT 算法，迁移到其他小规模数据集进行训练，与使用 CNN 结构的SOTA算法精度对比。

图2：ViT模型精度对比

图中前3列为不同尺度的ViT模型，使用不同的大规模数据集进行预训练，并迁移到各个子任务上的结果。第4列为BiT算法基于JFT-300M数据集预训练后，迁移到各个子任务的结果。第5列为2020年提出的半监督算法 Noisy Student 在 ImageNet 和 ImageNet ReaL 数据集上的结果。

---

说明：

BiT 与 Noisy Student 均为2020年提出的 SOTA 算法。

BiT算法：使用大规模数据集 JFT-300M 对 ResNet 结构进行预训练，其中，作者发现模型越大，预训练效果越好，最终指标最高的为4倍宽、152层深的 $ResNet152\times 4$。论文地址：Big Transfer (BiT): General Visual Representation Learning

Noisy Student 算法：使用知识蒸馏的技术，基于 EfficientNet 结构，利用未标签数据，提高训练精度。论文地址：Self-training with Noisy Student improves ImageNet classification

---

接下来，分别看一下ViT算法的各个组成部分。

图像分块嵌入

考虑到之前课程中学习的，Transformer结构中，输入需要是一个二维的矩阵，矩阵的形状可以表示为 $(N,D)$，其中 $N$ 是sequence的长度，而 $D$ 是sequence中每个向量的维度。因此，在ViT算法中，首先需要设法将 $H\times W\times C$ 的三维图像转化为 $(N,D)$ 的二维输入。

ViT中的具体实现方式为：将 $H\times W\times C$ 的图像，变为一个 $N\times (P^2\ast C)$ 的序列。这个序列可以看作是一系列展平的图像块，也就是将图像切分成小块后，再将其展平。该序列中一共包含了 $N=HW/P^2$ 个图像块，每个图像块的维度则是 $(P^2\ast C)$。其中 $P$ 是图像块的大小，$C$ 是通道数量。经过如上变换，就可以将 $N$ 视为sequence的长度了。

但是，此时每个图像块的维度是 $(P^2\ast C)$，而我们实际需要的向量维度是 $D$，因此我们还需要对图像块进行 Embedding。这里 Embedding 的方式非常简单，只需要对每个 $(P^2\ast C)$ 的图像块做一个线性变换，将维度压缩为 $D$ 即可。

上述对图像进行分块以及 Embedding 的具体方式如 图3 所示。

图3：图像分块嵌入示意图

具体代码实现如下所示。其中，使用了大小为 $P$ 的卷积来代替对每个大小为 $P$ 图像块展平后使用全连接进行运算的过程。

# coding=utf-8
# 导入环境
import os
import numpy as np
import cv2
from PIL import Image
import paddle
from paddle.io import Dataset
from paddle.nn import Conv2D, MaxPool2D, Linear, Dropout, BatchNorm, AdaptiveAvgPool2D, AvgPool2D
import paddle.nn.functional as F
import paddle.nn as nn

# 图像分块、Embedding
class PatchEmbed(nn.Layer):
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        # 原始大小为int，转为tuple，即：img_size原始输入224，变换后为[224,224]
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        # 图像块的个数
        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
        # kernel_size=块大小，即每个块输出一个值，类似每个块展平后使用相同的全连接层进行处理
        # 输入维度为3，输出维度为块向量长度
        # 与原文中：分块、展平、全连接降维保持一致
        # 输出为[B, C, H, W]
        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
        assert H == self.img_size[0] and W == self.img_size[1], \
            "Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
        # [B, C, H, W] -> [B, C, H*W] ->[B, H*W, C]
        x = self.proj(x).flatten(2).transpose((0, 2, 1))
        return x

Multi-head Attention

将图像转化为 $N\times (P^2\ast C)$ 的序列后，就可以将其输入到 Tranformer 结构中进行特征提取了。 Tranformer 结构中最重要的结构就是 Multi-head Attention，即多头注意力结构，如 图4 所示。

图4：Multi-head Attention 示意图

具有2个head的 Multi-head Attention 结构如 图5 所示。输入 $a^i$ 经过转移矩阵，并切分生成 $q^{(i,1)}$、$q^{(i,2)}$、$k^{(i,1}$