Vision Transformer

Vision Transformer

本文关注ViT论文4.5 Inspecting Vision Transformer可视化的原理及实现，此外还对ViT pytorch源码实现进行理解

Introduction

Title

An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

一张图像就相当于一些16x16的单词：将transformer应用到大规模图像识别任务上

transformer是2017年提出基于自注意力机制解决翻译任务的方法，一经推出就在nlp（neutral language process）领域爆火，image recognition是cv（computer vision）领域的基本任务之一。本文将nlp领域的transformer方法应用到cv领域的图像识别任务上，故将其称为vision transformer，简称ViT

Usage

ViT自2020年提出之后，迅速在cv领域爆火，论文引用量已超18k（2023.7）图片来源

ViT甚至在2022年超越了经典的ResNet架构的CNN模型，成为图像分类任务中的主流方法图片来源

目前ViT主要应用在图像分类、语义分割、目标检测这三个主流的cv任务上图片来源

目前ImageNet图像分类数据集上表现最好的3个方法中，都是基于transformer的结构，而ViT是将transformer应用到cv领域的开山之作，所以有必要认真掌握，就像掌握AlexNet和ResNet那样图片来源

Method

Overview

ViT 模型动态图展示了ViT模型的整体结构和forward过程，学习过transformer之后，再学习这篇文章，相对就会感觉比较轻松。

ViT的主要思想就是将transformer模型中Encoder部分直接搬过来，想办法将其应用到图片上。就ViT注意力机制部分来说，比transformer还要简单，只有Encoder，没有Decoder，以及Encoder-Decoder Cross Attention。

transformer用来处理seq2seq的任务，也就是输入和输出都是序列，Encoder的输入和输出也都是序列。那么要将其应用到图像分类任务上，关键是解决2个问题

如何将图片转换成Encoder能够接受的输入，也就是将图片转换为序列（tokens）
如何利用Encoder的输出来做分类任务，最终需要输出的图像的类别分布

基于上述关键问题，我将模型拆解为3个部分进行理解，其中Encoder部分在transformer中已经学习过，所有重点是Input和Output部分，随着学习的深入会发现，Output部分非常简单，就是接上一个MLP的分类头即可，所以重中之重是理解Input部分。

Input

将原始图像转换为Encoder能够接受的输入
Encoder

transformer原封不动搬过来
Output

将经过Encoder处理的输出用来做分类

论文中4个公式分别与3个部分的对应关系如下

Input

基本思想是将原始图片拆分成多个小图（对应了论文的标题），再把这些小图变换成Encoder的输入

这部分由可拆解成4部分来理解

将原始图像拆分成固定大小的小图片，论文中称之为patch，并展成向量
对patches向量做线性映射，映射到特定的长度，也就是Encoder输入token的维度
拼接上class token，此方法参考了BERT
对每个token（包括class token）加上可学习的位置编码

Encoder

参见transformer学习：

https://jalammar.github.io/illustrated-transformer/
transformer论文《Attention is all you need》，arXiv:1706.03762
台大李宏毅视频课程：https://www.bilibili.com/video/BV1Wv411h7kN

Output

使用class token经过Encoder后的输出向量，接上一个MLP分类头，映射到类别数量即可

Experiments

ViT论文根据提出了大小不同的3个模型

实验结果如下图

通过比较可以得出结论

更大的ViT模型有更好的表现
ViT模型只有在大规模数据集上预训练后，准确度才能超过传统的ResNet架构模型
ViT模型的训练开销要比传统模型小

Explainability

可视化部分共4个图，是老板提出要重点理解和实现的内容，分别是

Linear Embedding过滤器学习的内容
Position Embedding
Attention Distance随网络深度增加的变化
ViT模型关注图像中的哪些内容

Visualize Filters of Linear Embedding

这一部分是花最长时间搞懂的部分，论文中并没有提到如何实现，在探索的过程中走了很多弯路，原因有2个

论文中模型描述和实现代码的差异
自身知识的局限

先描述一下我在探索的过程中遇到的问题

论文中提到，要可视化Linear Embedding层中Filters学习到的东西，我当时就陷入了矛盾，可视化filters是可视化卷积层的filters，而我们现在是线性层啊，怎么可视化线性层呢？于是就想去找可视化线性层的相关资料，找来找去都是可视化卷积层的。后来又去找可视化卷积层的相关资料，发现入了一个大坑，有很多方法，什么可视化特征图、反卷积，而且一开始只想到用在python中使用matlab的库绘图，本身对这些工具的使用就不熟悉，导致很郁闷。

后来又去仔细看了一下代码，发现：Linear Embedding层是用卷积核大小和步长都为patch_size，输出通道数为768的卷积层实现的！此时，虽然还不知道如何实现，但最起码有了方向，去查如何可视化卷积层就行了

查相关资料的时候以结果为导向，直接看博客中有没有和论文中相似的可视化图，很快就查到这篇博客，看到其中贴的图和论文中的图特别相似，并且有代码实现，使用tensorboard实现的可视化，于是把代码拿过来改了一下，实现了和论文中类似的可视化结果

修改后的代码可视化Linear Embedding层卷积层的权重，代码如下

import os
import torch
import numpy as np
from torch.utils.tensorboard import SummaryWriter
import torchvision.utils as vutils
from modeling import VisionTransformer, CONFIGS
from sklearn.decomposition import PCA
from einops import rearrange

# 环境设置
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
log_dir = os.path.join(BASE_DIR, "results")
writer = SummaryWriter(log_dir=log_dir, filename_suffix="_kernel")
imagenet_labels = dict(enumerate(open('attention_data/ilsvrc2012_wordnet_lemmas.txt')))

# 模型选择：vit Large patch_size=32
config = CONFIGS["ViT-L_32"]
# 创建模型 并加载预训练权重
model = VisionTransformer(config, num_classes=21843, zero_head=False, img_size=224, vis=True)
model.load_from(np.load("attention_data/ViT-L_32.npz"))
model.eval()

# 取卷积层权重 X:(1024, 3, 32, 32)
X = model.transformer.embeddings.patch_embeddings.weight.detach()
# 将后3个维度合并，使用主成分分析工具要求输入维度为2
X = rearrange(X, 'n c h w -> n (c h w)')
# 与ViT论文一致，做28个主成分; pca components: (28, 3072)
pca = PCA(n_components=28)
pca.fit(X)
# 再将主成分转换成tensorboard 能够展示的RGB的形式
# filters: (28, 3, 32, 32)
filters = torch.tensor(pca.components_.reshape((-1, 3, 32, 32)))
# 使用tensorboard绘制filters的RGB图像
grid = vutils.make_grid(filters, normalize=True, scale_each=True, nrow=7)
writer.add_image("RGB embedding filters\nfirst 28 principal components", grid, global_step=620)

Visualize Positon Embedding

对于196个position embedding，计算每一个position embedding和其他position embedding之间的余弦相似度，会得到一个矩阵，矩阵的第 i 行表示，第 i 个position embedding和其他position embedding的相关性，将每一行转换成（14，14）小图，然后绘制出来即可，总共有14x14个这样的小图，每一个小图表示该位置上的patch和其他patch的相关性

实现代码如下

import os
import torch
import numpy as np
import matplotlib.pyplot as plt
from tqdm import tqdm
from modeling import VisionTransformer, CONFIGS
from einops import rearrange

# 环境设置
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
imagenet_labels = dict(enumerate(open('attention_data/ilsvrc2012_wordnet_lemmas.txt')))


# 定义计算两个向量余弦相似度函数
def cosine_similarity(x, y, norm=False):
    """ 计算两个向量x和y的余弦相似度 """
    assert len(x) == len(y), "len(x) != len(y)"

    xy = x.dot(y)
    x2y2 = np.linalg.norm(x, ord=2) * np.linalg.norm(x, ord=2)
    sim = xy / x2y2
    return sim


# 模型选择：vit base
config = CONFIGS["ViT-B_16"]
# 创建模型 并加载预训练模型权重
model = VisionTransformer(config, num_classes=1000, zero_head=False, img_size=224, vis=True)
model.load_from(np.load("attention_data/ViT-B_16-224.npz"))
model.eval()

# 获取可学习的位置编码
pos = model.transformer.embeddings.position_embeddings  # pos: (1, 197, 768)
# 移除维度大小为1的维度，且忽略[cls] token，并转换为numpy类型
pos = torch.squeeze(pos)[1:, :].detach().numpy()  # pos: numpy (196, 768)

# 计算similarity matrix
result = np.zeros((196, 196))
for i in tqdm(range(196)):
    for j in range(196):
        result[i, j] = cosine_similarity(pos[i, :], pos[j, :])
# result: (196, 196) => (196, 14, 14)
result = rearrange(result, 'n (h w) -> n h w', h=14, w=14)

# 绘制图像
fig, axs = plt.subplots(nrows=14, ncols=14, figsize=(14, 14),
                        subplot_kw={'xticks': [], 'yticks': []})
i = 0
for ax in axs.flat:
    ax.imshow(result[i, :, :], cmap='viridis')
    i += 1
plt.tight_layout()
plt.show()

Visualize Attention Distance with Network Depth

还没有看对应的实现，只是理解了一下。attention distance就相当于是卷积神经网络中的卷积层的感受野

在网络的较低层，由于注意力机制的存在，模型既关注一些小的区域，也关注一些大的区域，随着网络深度的增加，模型越来越关注更高的抽象层次，也就是图片的语义部分，感受野就比较大（个人感觉我这段理解写的很垃圾）

Visualize Attention Map

ViT在大规模数据集上与训练之后，在图像分类任务上取得了非常好的表现。想要深入了解模型，探索模型的可解释性，就要搞清模型根据图像的哪些部分做出的决策，或者说模型在这个任务上关注了图像的哪些部分。将模型关注的图像的部分可视化，也就是绘制特征图（Attention Map）

论文中明确提到了使用Attention Rollout方法绘制特征图，这是另外一篇论文提出的可视化transformer注意力流的其中一个方法

Samira Abnar and Willem Zuidema. Quantifying attention flow in transformers. In ACL, 2020

arXiv:2005.00928

在附录中作者也简要介绍了这个方法的思想，将ViT模型每一个注意力层的多个头的权重平均一下，然后把所有注意力层的多头平均权重累乘，得到的结果就能够表示整个模型的注意力

Briefly, we averaged attention weights of ViTL/16 across all heads and then recursively multiplied the weight matrices of all layers. This accounts for the mixing of attention across tokens through all layers.

为了更好的理解这个方法，首先要明确一个概念，注意力权重（Attention Weights）或称注意力矩阵图片来源

熟悉transformer的可以非常好的理解这个概念，也就是每个token和其他token计算的attention score所构成的矩阵

Attention Rollout方法的说明如下，简单理解就是将所有Attention Weight累乘，但是这里面要考虑多头和残差连接的处理

具体实现如下图所示，这里的Attention Maxtrix是运行过attention rollout算法之后的结果，通过reshape，再上采样之后得到一个类似于掩码功能的图像mask，再将mask作用到原图像上可视化出来即可

至于HeatMap的生成方式，同样也是使用mask作用在原始图像上，只不过使用python相关包对mask进行以下处理即可

实现代码如下代码参考

import os

import torch
import numpy as np
import cv2
from urllib.request import urlretrieve
from PIL import Image
from torchvision import transforms
from modeling import VisionTransformer, CONFIGS

# 环境设置
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
imagenet_labels = dict(enumerate(open('ilsvrc2012_wordnet_lemmas.txt')))

# Prepare Model
# 模型选择：vit base
config = CONFIGS["ViT-B_16"]
# 创建模型
model = VisionTransformer(config, num_classes=1000, zero_head=False, img_size=224, vis=True)
# 加载预训练模型权重
model.load_from(np.load("ViT-B_16-224.npz"))
model.eval()

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
im = Image.open("attention_data/plane2.png")
x = transform(im)
print(x.size())

# inference
logits, att_mat = model(x.unsqueeze(0))  # att_mat (list:12) elem:(1, 12, 197, 197)
att_mat = torch.stack(att_mat)  # att_mat (12, 1, 12, 197, 197)
att_mat = att_mat.squeeze(1)  # att_mat (12, 12, 197, 197)

# Average the attention weights across all heads.
att_mat = torch.mean(att_mat, dim=1)  # att_mat (12, 12, 197, 197) => (12, 197, 197)

# To account for residual connections, we add an identity matrix to the
# attention matrix and re-normalize the weights.
residual_att = torch.eye(att_mat.size(1))
aug_att_mat = att_mat + residual_att
aug_att_mat = aug_att_mat / aug_att_mat.sum(dim=-1).unsqueeze(-1)

# Recursively multiply the weight matrices	
joint_attentions = torch.zeros(aug_att_mat.size())
joint_attentions[0] = aug_att_mat[0]
for n in range(1, aug_att_mat.size(0)):
    joint_attentions[n] = torch.matmul(aug_att_mat[n], joint_attentions[n - 1])

# Attention from the output token to the input space.
v = joint_attentions[-1]  # v: (197, 197)
grid_size = int(np.sqrt(aug_att_mat.size(-1)))  # grid_size: 14
mask = v[0, 1:].reshape(grid_size, grid_size).detach().numpy()  # mask: (14, 14)
mask = cv2.resize(mask / mask.max(), im.size)[..., np.newaxis]  # mask: (224, 224, 1)
# 将mask作用到原始图片上
result = (mask * im).astype("uint8")
# 以下为绘图操作
fig, (ax1, ax2) = plt.subplots(ncols=2, figsize=(16, 16))
ax1.set_title('Original')
ax2.set_title('Attention Map')
ax1.imshow(im)
ax2.imshow(result)
plt.show()


# 以下为实现heatmap的代码 

# def show_mask_on_image(img, mask):
#     img = np.float32(img) / 255
#     heatmap = cv2.applyColorMap(np.uint8(255 * mask), cv2.COLORMAP_JET)
#     heatmap = np.float32(heatmap) / 255
#     cam = heatmap + np.float32(img)
#     cam = cam / np.max(cam)
#     return np.uint8(255 * cam)
# 
# 
# np_img = np.array(im)[:, :, ::-1]
# mask = show_mask_on_image(np_img, mask)
# cv2.imshow("Input Image", np_img)
# cv2.imshow("Output Image", mask)
# cv2.waitKey(-1)

Code

代码参考

代码结构框架如下图所示

Model Variants

以ViT Base-16模型作为例子理解

# Prepare Model
# 模型选择：vit base
config = CONFIGS["ViT-B_16"]
# 创建模型
model = VisionTransformer(config, num_classes=1000, zero_head=False, img_size=224, vis=True)
# 加载预训练模型权重
model.load_from(np.load("attention_data/ViT-B_16-224.npz"))
model.eval()

配置参数如下

def get_b16_config():
    """Returns the ViT-B/16 configuration."""
    config = ml_collections.ConfigDict()
    config.patches = ml_collections.ConfigDict({'size': (16, 16)})
    config.hidden_size = 768
    config.transformer = ml_collections.ConfigDict()
    config.transformer.mlp_dim = 3072
    config.transformer.num_heads = 12
    config.transformer.num_layers = 12
    config.transformer.attention_dropout_rate = 0.0
    config.transformer.dropout_rate = 0.1
    config.classifier = 'token'
    config.representation_size = None
    return config

阅读类代码时，主要关注forward函数，根据需要查看__init__函数中的定义，搞清类别之间的调用关系，对照模型架构图进行理解，补充学习代码细节实现（pytorch）