手写VIT

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, repeat
from einops.layers.torch import Rearrange

class pre_process(nn.Module):
    def __init__(self, image_size, patch_size, patch_dim, dim):
        """
        image_size:输入图像的大小,例如224
        patch_size:每个patch的大小,例如16
        patch_dim: 每个patch的特征维度,例如768
        dim: 模型的隐藏层维度,例如768
        """
        super().__init__()
        self.patch_size = patch_size
        self.dim = dim
        # 计算图像中patch的数量，并赋值给类属性 self.patch_num
        self.patch_num = (image_size // patch_size) ** 2
        # 初始化一个线性层 self.linear_embedding,用于将每个patch的特征维度从patch_dim映射到dim
        self.linear_embedding = nn.Linear(patch_dim, dim)
        self.position_embedding = nn.Parameter(torch.randn(1,self.patch_num+1, self.dim)) # 使用广播
        self.CLS_token = nn.Parameter(torch.randn(1, 1, self.dim)) # 别忘了维度要和(B,L,C)对齐

    def forward(self, x):
        """
        输入的x的数组表示为(B,C,H,W)，需要将它划分为(B,L,C)
        输入形状- B:批次大小  C:通道数  H:高度 W:宽度
        输出形状- B:批次大小  L:patch的个数,等于(H/p1)*(W/p2)  C:每个patch的像素值数量,等于 p1*p2*channels
        """
        x = rearrange(x, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = self.patch_size, p2 = self.patch_size) # (B, L, C)
        """

        """
        x = self.linear_embedding(x)
        b, l, c = x.shape
        # 复制后的CLS_token的形状为(b, 1, d),表示每个样本都拥有一个CLS Token
        CLS_token = repeat(self.CLS_token, '1 1 d -> b 1 d', b=b) # 位置编码复制 B 份
        # 级联CLS Token
        x = torch.concat((CLS_token, x), dim=1) 
        # 位置嵌入
        # self.position_embedding的形状是(1, patch_num+1, dim)，x的形状是(batch_num, patch_length+1, dim).
        # 在进行加法时,self.position_embedding会被扩展到(batch_size, patch_num+1, dim),从而与x进行逐元素相加
        x = x + self.position_embedding #batch的数量没有发生变化，但是batch的

class Multihead_self_attention(nn.Module):
    def __init__(self, heads, head_dim, dim):
        super().__init__()
        """
        self.dim: 输入数据的维度  
        self.head_dim: 每个注意力头的维度
        self.heads: 注意力头的数量
        self.inner_dim: 多头自注意力最后的输出维度，等于 heads*head_dim
        """
        self.head_dim = head_dim # 每一个注意力头的维度
        self.heads = heads # 注意力头个数
        self.inner_dim = self.heads*self.head_dim #多头注意力最后输出维度
        self.scale = self.head_dim**-0.5 #正则化系数
        # 一个线性层，用于将输入数据映射到查询(Q)，键(K)和值(V)三个矩阵，每个矩阵的维度为inner_dim
        self.to_qkv = nn.Linear(dim, self.inner_dim*3) # 生成qkv,每一个矩阵的维度由自注意力头的维度以及头的个数决定
        # 一个线性层，用于将多头自注意力输出结果映射回原始维度dim
        self.to_output = nn.Linear(self.inner_dim, dim)
        # 一个层归一化层，用于对输入数据进行归一化
        self.norm = nn.LayerNorm(dim)
        # 一个Softmax函数，用于对注意力得分进行归一化
        self.softmax = nn.Softmax(dim=-1)
    
    def forward(self, x):
        x = self.norm(x) # PreNorm
        # .chunk()方法是PyTorch中用于分割张量的函数；dim=-1表示沿着最后一个维度进行分割
        qkv = self.to_qkv(X).chunk(3, dim=-1) # 划分QKV, 返回一个列表， 其中就包含了QKV
        # 对 QKV 的多头映射进行拆分， 得到(B, head, L, head_dim)
        """
        map函数会对qkv列表中的每个张量应用lambda表达式
        'b l (h dim) -> b h l dim'表示将输入张量的维度从(batch_size, sequence_length, num_heads*head_dim)变换为(batch_size, num_heads, sequence_length, head_dim)
        dim=self.head_dim 是指定head_dim 的值，以便正确拆分
        """
        Q, K, V = map(lambda t: rearrange(t, 'b l (h dim) -> b h l dim', dim=self.head_dim), qkv)
        K_T = K.transpose(-1, -2) # 对K进行转置，用于计算自注意力
        att_score = Q@K_T*self.scale # 计算自注意力得分
        att = self.softmax(att_score) # softmax
        out = att@V # (B, H, L, dim):自注意力输出
        # 将输出张量的维度从(batch_size, num_heads, sequence_length, head_dim)变换为 (batch_size, sequence_length, num_heads*head_dim)
        out = rearrange(out, 'b h l dim -> b l (h dim)') # 拼接
        output = self.to_output(out) # 输出映射
        return output

class FeedForward(nn.Module):
    def __init__(self, dim, mlp_dim):
        super().__init__()
        self.fc1 = nn.Linear(dim, mlp_dim)
        self.fc2 = nn.Linear(mlp_dim, dim)
        self.norm = nn.LayerNorm(dim)
    
    def forward(self, x):
        x = self.norm(x)
        x = F.gelu(self.fc1(x))
        x = self.fc2(x)
        return x

class Transformer_block(nn.Module):
    def __init__(self, dim, heads, head_dim, mlp_dim, depth, num_class):
        super().__init__()
        self.to_patch_embedding = pre_process(image_size=image_size, patch_size=patch_size)