0
0

【计算机视觉】ViT：代码逐行解读

文章目录

一、代码

  1. import torch
  2. import torch.nn as nn
  3. from einops import rearrange
  5. from self_attention_cv import TransformerEncoder
  7. classViT(nn.Module):def__init__(self,*,
  8.                  img_dim,
  9.                  in_channels=3,
  10.                  patch_dim=16,
  11.                  num_classes=10,
  12.                  dim=512,
  13.                  blocks=6,
  14.                  heads=4,
  15.                  dim_linear_block=1024,
  16.                  dim_head=None,
  17.                  dropout=0, transformer=None, classification=True):"""
  18.         Args:
  19.             img_dim: the spatial image size
  20.             in_channels: number of img channels
  21.             patch_dim: desired patch dim
  22.             num_classes: classification task classes
  23.             dim: the linear layer's dim to project the patches for MHSA
  24.             blocks: number of transformer blocks
  25.             heads: number of heads
  26.             dim_linear_block: inner dim of the transformer linear block
  27.             dim_head: dim head in case you want to define it. defaults to dim/heads
  28.             dropout: for pos emb and transformer
  29.             transformer: in case you want to provide another transformer implementation
  30.             classification: creates an extra CLS token
  31.         """super().__init__()assert img_dim % patch_dim ==0,f'patch size {patch_dim} not divisible'
  32.         self.p = patch_dim
  33.         self.classification = classification
  34.         tokens =(img_dim // patch_dim)**2
  35.         self.token_dim = in_channels *(patch_dim **2)
  36.         self.dim = dim
  37.         self.dim_head =(int(dim / heads))if dim_head isNoneelse dim_head
  38.         self.project_patches = nn.Linear(self.token_dim, dim)
  40.         self.emb_dropout = nn.Dropout(dropout)if self.classification:
  41.             self.cls_token = nn.Parameter(torch.randn(1,1, dim))
  42.             self.pos_emb1D = nn.Parameter(torch.randn(tokens +1, dim))
  43.             self.mlp_head = nn.Linear(dim, num_classes)else:
  44.             self.pos_emb1D = nn.Parameter(torch.randn(tokens, dim))if transformer isNone:
  45.             self.transformer = TransformerEncoder(dim, blocks=blocks, heads=heads,
  46.                                                   dim_head=self.dim_head,
  47.                                                   dim_linear_block=dim_linear_block,
  48.                                                   dropout=dropout)else:
  49.             self.transformer = transformer
  51.     defexpand_cls_to_batch(self, batch):"""
  52.         Args:
  53.             batch: batch size
  54.         Returns: cls token expanded to the batch size
  55.         """return self.cls_token.expand([batch,-1,-1])defforward(self, img, mask=None):
  56.         batch_size = img.shape[0]
  57.         img_patches = rearrange(
  58.             img,'b c (patch_x x) (patch_y y) -> b (x y) (patch_x patch_y c)',
  59.                                 patch_x=self.p, patch_y=self.p)# project patches with linear layer + add pos emb
  60.         img_patches = self.project_patches(img_patches)if self.classification:
  61.             img_patches = torch.cat((self.expand_cls_to_batch(batch_size), img_patches), dim=1)
  63.         patch_embeddings = self.emb_dropout(img_patches + self.pos_emb1D)# feed patch_embeddings and output of transformer. shape: [batch, tokens, dim]
  64.         y = self.transformer(patch_embeddings, mask)if self.classification:# we index only the cls token for classification. nlp tricks :Preturn self.mlp_head(y[:,0,:])else:return y

二、代码解读

2.1 大体理解

这段代码是一个实现了 Vision Transformer（ViT）模型的 PyTorch 实现。

ViT 是一个基于 Transformer 架构的图像分类模型，其主要思想是将图像分成一个个固定大小的 patch ，并将这些 patch 看做是一个个 token 输入到 Transformer 中进行特征提取和分类。

以下是对代码的解读：

  1. ViT类继承自nn.Module类，其构造函数有一系列参数，包括输入图像的尺寸、patch的大小、输出类别数、注意力机制中的头数等等。
  2. project_patches函数通过一个全连接层将每个patch映射到一个d维的特征空间中。
  3. 如果classification = True，则将一个额外的CLS token添加到输入的token序列的开头，即对于每张图像添加一个形状为[1, 1, d]的CLS token。同时，在ViT中采用的是绝对位置编码，因此还添加了一个1D的位置编码向量，其形状为[num_patches + 1, d]，其中num_patches表示图像被划分成的patch数目。如果classification = False，则不添加CLS token。
  4. forward函数首先将输入的图像进行patch划分，并通过project_patches函数将每个patch映射到d维特征空间中。接着，将位置编码向量加到映射后的patch特征向量上，并进行dropout处理。如果classification=True，则在特征序列开头添加CLS token。接着将这些特征输入到Transformer中，进行特征提取。最后输出分类结果，如果classification=True，则只返回CLS token的分类结果。

2.2 详细理解

  1. from self_attention_cv import TransformerEncoder
  1. self_attention_cv

是一个基于PyTorch实现的库，提供了在计算机视觉任务中使用自注意力机制的模块和网络，例如

  1. Transformer Encoder

和

  1. Attention Modules

。

它主要针对图像分类、对象检测、语义分割等任务，支持多种自注意力模块的实现，包括

  1. Simplified Self-Attention

、

  1. Full Self-Attention

和

  1. Local Self-Attention

等。此外，该库还提供了一些常见的计算机视觉任务模型的实现，例如

  1. Vision Transformer（ViT）

和

  1. Swin Transformer

等。

  1. TransformerEncoder

是一个自注意力机制的编码器，用于将输入序列转换为编码后的序列。自注意力机制允许模型能够根据输入序列中的其他位置来加权计算每个位置的表示。这种机制在自然语言处理中的应用非常广泛，比如BERT、GPT等模型都采用了自注意力机制。

  1. TransformerEncoder

是基于PyTorch实现的，可以在计算机视觉任务中使用，例如图像分类、对象检测、语义分割等。它支持多头注意力、残差连接和LayerNorm等特性。在这个代码中，ViT模型中的Transformer部分采用了TransformerEncoder作为默认的实现。

  1. def__init__(self,*,
  2.                 img_dim,
  3.                 in_channels=3,
  4.                 patch_dim=16,
  5.                 num_classes=10,
  6.                 dim=512,
  7.                 blocks=6,
  8.                 heads=4,
  9.                 dim_linear_block=1024,
  10.                 dim_head=None,
  11.                 dropout=0, transformer=None, classification=True):super().__init__()assert img_dim % patch_dim ==0,f'patch size {patch_dim} not divisible'
  12.     self.p = patch_dim
  13.     self.classification = classification
  14.     tokens =(img_dim // patch_dim)**2
  15.     self.token_dim = in_channels *(patch_dim **2)
  16.     self.dim = dim
  17.     self.dim_head =(int(dim / heads))if dim_head isNoneelse dim_head
  18.     self.project_patches = nn.Linear(self.token_dim, dim)
  20.     self.emb_dropout = nn.Dropout(dropout)if self.classification:
  21.         self.cls_token = nn.Parameter(torch.randn(1,1, dim))
  22.         self.pos_emb1D = nn.Parameter(torch.randn(tokens +1, dim))
  23.         self.mlp_head = nn.Linear(dim, num_classes)else:
  24.         self.pos_emb1D = nn.Parameter(torch.randn(tokens, dim))if transformer isNone:
  25.         self.transformer = TransformerEncoder(dim, blocks=blocks, heads=heads,
  26.                                                 dim_head=self.dim_head,
  27.                                                 dim_linear_block=dim_linear_block,
  28.                                                 dropout=dropout)else:
  29.         self.transformer = transformer

这段代码定义了一个名为 ViT 的 PyTorch 模型类，它是一个使用自注意力机制（Self-Attention）实现的视觉 Transformer 模型。其中主要参数包括：

  • img_dim：输入图片的空间大小
  • in_channels：输入图片的通道数
  • patch_dim：将图片划分成固定大小的 patch 的大小
  • num_classes：分类任务的类别数
  • dim：线性层的维度，用于将每个 patch 投影到 MHSA 空间
  • blocks：Transformer 模型中的块数
  • heads：注意力头的数量
  • dim_linear_block：线性块内部的维度
  • dim_head：每个头的维度，如果没有指定则默认为 dim/heads
  • dropout：用于位置编码和 Transformer 的 dropout 概率
  • transformer：可选的 TransformerEncoder 类实例
  • classification：是否包含额外的 CLS 标记以用于分类任务
  1. def__init__(self,*,
  2.                  img_dim,
  3.                  in_channels=3,
  4.                  patch_dim=16,
  5.                  num_classes=10,
  6.                  dim=512,
  7.                  blocks=6,
  8.                  heads=4,
  9.                  dim_linear_block=1024,
  10.                  dim_head=None,
  11.                  dropout=0, transformer=None, classification=True):super().__init__()

这里定义了 ViT 类的构造函数，其包含多个参数，包括输入图像大小 

  1. img_dim

，输入通道数 

  1. in_channels

，分块大小 

  1. patch_dim

，分类数目 

  1. num_classes

，嵌入维度 

  1. dim

，Transformer编码器的块数 

  1. blocks

，头数 

  1. heads

，线性块的维度 

  1. dim_linear_block

，注意力头维度 

  1. dim_head

，Dropout概率 

  1. dropout

，可选的Transformer编码器 

  1. transformer

，以及是否进行分类的标志 

  1. classification

。

  1. assert img_dim % patch_dim ==0,f'patch size {patch_dim} not divisible'
  2.     self.p = patch_dim
  3.     self.classification = classification

这里检查 

  1. img_dim

是否能够被 

  1. patch_dim

整除，如果不能整除，则会引发断言错误。同时，将 

  1. patch_dim

存储到 

  1. self.p

中，并将是否进行分类的标志存储到 

  1. self.classification

中。

  1.     tokens =(img_dim // patch_dim)**2
  2.     self.token_dim = in_channels *(patch_dim **2)
  3.     self.dim = dim
  4.     self.dim_head =(int(dim / heads))if dim_head isNoneelse dim_head
  5.     self.project_patches = nn.Linear(self.token_dim, dim)

这里计算了输入图像中可分块的数量 

  1. tokens

，并将每个块的维度 

  1. self.token_dim

设置为 

  1. in_channels * (patch_dim ** 2)

。

将嵌入维度 

  1. dim

存储到 

  1. self.dim

中，并根据 

  1. dim_head

是否为 None，设置注意力头维度 

  1. self.dim_head

。

  1. self.project_patches

是一个线性层，用于将每个块投影到嵌入空间中。

  1.     self.emb_dropout = nn.Dropout(dropout)if self.classification:
  2.         self.cls_token = nn.Parameter(torch.randn(1,1, dim))
  3.         self.pos_emb1D = nn.Parameter(torch.randn(tokens +1, dim))
  4.         self.mlp_head = nn.Linear(dim, num_classes)else:
  5.         self.pos_emb1D = nn.Parameter(torch.randn(tokens, dim))

这里定义了嵌入层的

  1. Dropout

层，并根据是否进行分类的标志，设置类别标记 

  1. self.cls_token

、位置嵌入 

  1. self.pos_emb1D

和 

  1. MLP

头 

  1. self.mlp_head

。如果不进行分类，则不需要 

  1. self.cls_token

和 

  1. self.mlp_head

。

  1. if transformer isNone:
  2.         self.transformer = TransformerEncoder(dim, blocks=blocks, heads=heads,
  3.                                                 dim_head=self.dim_head,
  4.                                                 dim_linear_block=dim_linear_block,
  5.                                                 dropout=dropout)else:
  6.         self.transformer = transformer
  1. self.emb_dropout = nn.Dropout(dropout):

定义了一个

  1. dropout

层，用于在embedding后对其进行dropout操作。

  1. if self.classification::

如果是分类任务，就执行下面的操作，否则跳过。

  1. self.cls_token = nn.Parameter(torch.randn(1, 1, dim)):

定义了一个可训练参数

  1. cls_token

，表示分类token，它是一个1x1xdim的tensor，其中dim表示embedding维度。

  1. self.pos_emb1D = nn.Parameter(torch.randn(tokens + 1, dim)):

定义了一个可训练参数pos_emb1D，表示位置嵌入，它是一个(tokens+1)xdim的tensor，其中tokens表示图像被分成的patch数，dim表示embedding维度。

  1. self.mlp_head = nn.Linear(dim, num_classes):

定义了一个全连接层，将embedding映射到输出类别的数量。

最后，根据传入的参数来选择使用默认的TransformerEncoder，还是使用传入的transformer。如果没有传入，则使用默认的TransformerEncoder，否则使用传入的transformer。

  1. defexpand_cls_to_batch(self, batch):"""
  2.     Args:
  3.         batch: batch size
  4.     Returns: cls token expanded to the batch size
  5.     """return self.cls_token.expand([batch,-1,-1])

该方法的作用是将 Transformer 中的分类 token 扩展到整个批次的样本数。它接受一个 batch 参数作为批次大小，返回一个形状为 [batch, 1, dim] 的张量，其中 dim 是 Transformer 模型的维度大小。在这个方法中，使用了 PyTorch 的 expand() 方法来实现扩展操作。

  1. defforward(self, img, mask=None):
  2.     batch_size = img.shape[0]
  3.     img_patches = rearrange(
  4.         img,'b c (patch_x x) (patch_y y) -> b (x y) (patch_x patch_y c)',
  5.                             patch_x=self.p, patch_y=self.p)# project patches with linear layer + add pos emb
  6.     img_patches = self.project_patches(img_patches)if self.classification:
  7.         img_patches = torch.cat((self.expand_cls_to_batch(batch_size), img_patches), dim=1)
  9.     patch_embeddings = self.emb_dropout(img_patches + self.pos_emb1D)# feed patch_embeddings and output of transformer. shape: [batch, tokens, dim]
  10.     y = self.transformer(patch_embeddings, mask)if self.classification:# we index only the cls token for classification. nlp tricks :Preturn self.mlp_head(y[:,0,:])else:return y

在 

  1. forward

函数中，接收输入的 

  1. img

和 

  1. mask

。

通过 

  1. img_dim

和 

  1. patch_dim

计算出 

  1. tokens

数量，其中 tokens 为图像分割成的块的数量。

将输入的 img 分成 patch，并通过 rearrange 函数重组成形状为 

  1. [batch_size, tokens, patch_dim * patch_dim * in_channels]

的张量。

通过 

  1. Linear

层将每个 patch 映射到 dim 维度，并加上位置编码 

  1. pos_emb1D

。

如果是用于分类任务，则在序列的开头插入一个 

  1. CLS token

，然后与处理后的 patch 张量按列拼接。

对 patch_embeddings 应用 dropout，并输入到 TransformerEncoder 中，返回输出张量 y，形状为 

  1. [batch_size, tokens, dim]

。

如果是用于分类任务，则从 y 中取出 CLS token，输入到一个 Linear 层中进行分类，输出分类结果。

如果不是分类任务，则直接返回 y。

标签： 计算机视觉 深度学习 人工智能
本文转载自: https://blog.csdn.net/wzk4869/article/details/130488137
版权归原作者 旅途中的宽~ 所有， 如有侵权，请联系我们删除。
登录后发布评论

“【计算机视觉】ViT：代码逐行解读”的评论:

还没有评论
关于作者
...
overfit同步小助手
文章同步
相关阅读

网络安全法-网络运行安全

使用selenium/drissionpage时如何阻止chrome自动跳转http到https

docker desktop 里部署的Open WebUI 管理员密码忘记了的处理方法

在ubuntu20.04中搭建onsite比赛运行环境

利用开源的低代码表单设计器FcDesigner高效管理和渲染复杂表单结构

Kafka学习笔记

【前端】浏览器输入url到页面呈现发生了什么？