ViT模型——pytorch实现

最新推荐文章于 2024-07-03 19:44:41 发布

CV_Peach

最新推荐文章于 2024-07-03 19:44:41 发布

阅读量2.7k

点赞数 7

文章标签： pytorch 深度学习自然语言处理

本文链接：https://blog.csdn.net/Peach_____/article/details/128767846

版权

论文传送门：AN IMAGE IS WORTH 16X16 WORDS: TRANSFORMERS FOR IMAGE RECOGNITION AT SCALE

ViT模型的特点：

抛弃Conv结构，使用Transformer对图像进行特征提取，完成图像分类任务。

ViT模型的结构：

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；

Attention计算公式
Multi-Head Attention计算公式
MLP：Linear + GELU + Dropout + Linear + Dropout；
③MLP Head：进行LayerNorm，提取class token，然后进行Linear，输出节点个数为类别数量。
(若针对ImageNet-21k数据集，则对class token进行Linear + Tanh + Linear)

不同规模的ViT模型：

不同规模的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