论文传送门:AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE
ViT模型的特点:
抛弃Conv结构,使用Transformer对图像进行特征提取,完成图像分类任务。
ViT模型的结构:
①Embedding:包括Patch Embedding、Position Embedding和Class Embedding;
Patch Embedding:将输入图像划分成一个又一个的子图像(Patch),并转换成向量序列(token);
Position Embedding:在Patch token前添加(拼接)class token,用于分类,为可训练参数;
Position Embedding:对token的位置信息进行编码,使用矩阵加法来实现,为可训练参数;
②Transformer Encoder:将Transformer Encoder Block重复堆叠L层,Transformer Encoder Block的结构如图右所示,为LayerNorm + Multi-Head Attention + Residual + LayerNorm + MLP + Residual;
LayerNorm:NLP领域常用的Normalization方法,计算公式与BN相似,但与Batch无关,而是对每个token(词)进行标准化处理,参考文献:Layer Normalization;
Multi-Head Attention:self-attention的一种,结构与计算公式如下,参考文献:Attention Is All You Need;
MLP:Linear + GELU + Dropout + Linear + Dropout;
③MLP Head:进行LayerNorm,提取class token,然后进行Linear,输出节点个数为类别数量。
(若针对ImageNet-21k数据集,则对class token进行Linear + Tanh + Linear)
不同规模的ViT模型:
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
class Embedding(nn.Module): # Patch Embedding + Position Embedding + Class Embedding
def __init__(self, image_channels=3, image_size=224, patch_size=16, dim=768, drop_ratio=0.):
super(Embedding, self).__init__()
self.num_patches = (image_size // patch_size) ** 2 # Patch数量
self.patch_conv = nn.Conv2d(image_channels, dim, patch_size, patch_size) # 使用卷积将图像划分成Patches
self.cls_token = nn.Parameter(torch.zeros(1, 1, dim)) # class embedding
self.pos_emb = nn.Parameter(torch.zeros(1, self.num_patches + 1, dim)) # position embedding
self.dropout = nn.Dropout(drop_ratio)
def forward(self, x):
x = self.patch_conv(x