[论文阅读]Vit（Vision Transformer）及代码

TRM

VIT 只用了encoder部分，就是把编码端用于图像

如何把图片放到encoder中：朴素思路，一个一个像素点输入

不直接输入的原因：224*224 参数数量太多，远大于nlp中的参数量

VIT

2部分中处理用的方法：1.把一个整体的patch拉成一个向量，再对其长度转变为embedding-size（拉长，减小） 2.Conv成768的channel数也行

3部分处理方法：

3.1 生成CLS符号（分类标签）的TokenEmbedding（*）【粉色的叫tokenEMB】

3.2 生成所有序列的位置编码（紫色的）【自注意力机制不含位置信息】

3.3 TokenEMB + 位置编码（紫色加粉色）

VIT模型

整体流程：

输入encoder中的数据格式 n（patch+1） * 768

最后5的步骤是，把第一个CLS拿出来，做一个多分类的任务，然后就能知道这张图片是属于哪个类的（可以认为CLS这一条向量，储存了类别信息）

也有不用CLS的做法，对每个token输出，并做一个GAP就能得到

一个 token 应该是一个形状为 (1, 768) 的向量（也就是一个 768 维的向量），而不是 (n, 768) 的张量。每个 token 代表了一个单独的特征向量，而不是多个 token 组成的张量。

下面是更正的步骤：

1. 汇总 Token：首先，您需要计算这些单个 token 的平均值，以获得一个形状为 (1, 768) 的全局特征向量。这可以通过在维度 0 上取平均值来实现。

pythonCopy code
global_feature = tokens.mean(dim=0)  # 计算全局平均特征向量

1. 添加分类头：接下来，您可以将全局特征向量传递给一个分类头，这个头通常是一个全连接层（Linear）层，输出维度为 1000（对应于 1000 个类别）。您可以使用 softmax 激活函数将输出转换为分类概率。

pythonCopy code
classifier_head = nn.Linear(768, 1000)  # 创建分类头，输入维度为 768，输出维度为 1000

1. 前向传播：在前向传播过程中，将全局特征向量传递给分类头，然后应用 softmax 函数以获取每个类别的概率分布。

pythonCopy codelogits = classifier_head(global_feature)  # 通过分类头获取 logits
probs = nn.functional.softmax(logits, dim=0)  # 使用 softmax 转换为概率分布

现在，probs 包含了每个类别的分类概率分布。可以使用这些概率来进行分类预测，通常选择概率最高的类别作为最终预测结果。

这就是如何使用 Global Average Pooling (GAP) 和分类头来进行 1000 类的分类预测。请确保在训练时使用适当的损失函数，如交叉熵损失，来训练模型以进行分类任务。

# 导入所需的库和模块
import torch
from torch import nn
from einops import rearrange, repeat
from einops.layers.torch import Rearrange

# 辅助函数，如果输入不是元组，则将其转换为元组
def pair(t):
    return t if isinstance(t, tuple) else (t, t)

# 定义 PreNorm 类，对输入进行预归一化，然后应用函数
class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)  # 使用 Layer Normalization 对输入进行归一化
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)  # 应用函数到归一化后的输入

# 定义 FeedForward 类，实现前馈神经网络
class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),  # 全连接层
            nn.GELU(),  # GELU激活函数
            nn.Dropout(dropout),  # Dropout层
            nn.Linear(hidden_dim, dim),  # 全连接层
            nn.Dropout(dropout)  # Dropout层
        )

    def forward(self, x):
        return self.net(x)  # 应用前馈神经网络到输入

# 定义 Attention 类，包括自注意力计算
class Attention(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.):
        super().__init__()
        inner_dim = dim_head * heads  # 内部维度计算 （聚合头特征）
        '''
        具体来说，dim_head 表示每个头的输出维度。假设输入序列的维度是 dim，
        并且自注意力机制被分为 num_heads 个头，那么每个头的输出维度就是 dim_head。
        因此，总的输出维度将是 num_heads * dim_head = inner_dim。
        一头Wk Wq Wv 变换成的维度为dim_head == 64
        '''

        project_out = not (heads == 1 and dim_head == dim)  # 是否需要投影输出

        self.heads = heads  # 多头注意力的头数
        self.scale = dim_head ** -0.5  # 缩放因子

        self.attend = nn.Softmax(dim=-1)  # 注意力权重计算使用Softmax
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)  # 生成查询、键、值的线性映射 *3主要是保证QKV的并行计算

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),  # 输出的全连接层
            nn.Dropout(dropout)  # 输出层的Dropout
        ) if project_out else nn.Identity()  # 如果不需要投影输出，使用恒等映射

    def forward(self, x):
        qkv = self.to_qkv(x).chunk(3, dim=-1)  # 对输入进行查询、键、值分割 x：1 196+1 1024 -> 1 197 512*3   qkv是一个元组 q k v 1 197 512
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)  # 重排分割后的查询、键、值
        #batch 序列长 （head*head_dims）-> batch head 序列长 head_dims(经过头后输出的维度)1 197 64*8 -> 1 8 197 64
        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale  # 计算点积注意力权重

        attn = self.attend(dots)  # 计算注意力分数 1 8 197 197 8个头，8个不同的注意力分数图(197*197，代表1与197个序列的相似度)

        out = torch.matmul(attn, v)  # 用注意力权重聚合值 1 8 197 64（1个头中， 每个位置融合了197个value后的新值）
        out = rearrange(out, 'b h n d -> b n (h d)')  # 重排输出1 197 8*64（融合为大的values）
        #得到聚合头的特征排列
        return self.to_out(out)  # 返回经过投影的输出

# 定义 Transformer 类，包括多层的自注意力和前馈神经网络
class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.layers = nn.ModuleList([])  # 存储Transformer层的列表
        for _ in range(depth): #堆叠trm  结合图看
            self.layers.append(nn.ModuleList([
                PreNorm(dim, Attention(dim, heads=heads, dim_head=dim_head, dropout=dropout)),  # 自注意力层
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))  # 前馈神经网络层mlp_dim是hidden_dim,出来的还是1024
            ]))

    def forward(self, x):
        for attn, ff in self.layers:  # 遍历每个Transformer层
            x = attn(x) + x  # 自注意力层和残差连接
            x = ff(x) + x  # 前馈神经网络层和残差连接
            print(x.shape)
        return x  # 返回Transformer层的输出

# 定义 ViT 类，包括块嵌入、位置嵌入、Transformer层、池化和多层感知机（MLP）
class ViT(nn.Module):
    def __init__(self, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool='cls', channels=3, dim_head=64, dropout=0., emb_dropout=0.):
        super().__init__()

        # 提取图像和块的尺寸信息
        image_height, image_width = pair(image_size)  #224*224
        patch_height, patch_width = pair(patch_size) #16*16

        assert image_height % patch_height == 0 and image_width % patch_width == 0, 'Image dimensions must be divisible by the patch size.'

        num_patches = (image_height // patch_height) * (image_width // patch_width)  # 计算块的数量
        patch_dim = channels * patch_height * patch_width  # 计算拉平向量的维度

        assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'

        #patch-embedding操作
        self.to_patch_embedding = nn.Sequential(
            Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1=patch_height, p2=patch_width),  # 通过重排操作将图像块flatten为向量 主要目的是将图像分割成许多小的块，然后将每个块中的像素值排列成一行，以便它们可以被输入到全连接层或其他神经网络层中
            nn.Linear(patch_dim, dim)  # 维度的变换   （h w）就是 h*w的意思 14*14 = 196（p1 p2 c）就是 p1 * p2 *c 的意思 
        )

        # 定义位置嵌入和类别标志
        self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))  # 初始化位置嵌入参数 （生成所有patches+cls的位置编码）【num-patch+1, dim】
        self.cls_token = nn.Parameter(torch.randn(1, 1, dim))  # 初始化类别标志参数 【1，dim】
        self.dropout = nn.Dropout(emb_dropout)  # 嵌入的Dropout层

        # 定义Transformer层
        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)

        # 池化方式和输出层
        self.pool = pool
        self.to_latent = nn.Identity()  # 恒等映射层

        self.mlp_head = nn.Sequential(
            nn.LayerNorm(dim),  # 使用Layer Normalization进行归一化
            nn.Linear(dim, num_classes)  # 输出层的多层感知机（MLP） 做多分类的任务
        )

    def forward(self, img):
        #想象成立方体来理解好理解一些
        # 对图像应用块嵌入
        x = self.to_patch_embedding(img)  # 1 3 224 224-> 1 14*14 1024 
        b, n, _ = x.shape  # 获取输入张量的形状

        # 重复类别标志并与块嵌入相连接
        cls_tokens = repeat(self.cls_token, '() n d -> b n d', b=b)  # 重复类别标志以匹配批次大小,每一个样本，都要加一个CLS
        print(cls_tokens.shape,'cls_token')
        x = torch.cat((cls_tokens, x), dim=1)  # 连接类别标志和块嵌入 在序列长度维度concat

        # 添加位置嵌入并应用dropout
        x += self.pos_embedding[:, :(n + 1)]  # 每个（197个）token【1，1024】加上位置嵌入
        x = self.dropout(x)  # 应用嵌入的Dropout层   batch 197 1024
        #print(x.shape)   #结果：batch 197 1024

        # 使用Transformer层进行特征提取
        x = self.transformer(x)

        # 池化层（平均池化或CLS池化）
        x = x.mean(dim=1) if self.pool == 'mean' else x[:, 0] #切片第0个元素
        print(x.shape)
        # 应用输出层的多层感知机（MLP）
        x = self.to_latent(x)
        return self.mlp_head(x)

# 创建ViT模型的实例，并设置参数
v = ViT(
    image_size=224,
    patch_size=16,
    num_classes=1000,  #CLS需要做多少类别的分类
    dim=1024, #维度
    depth=6, #encoder堆叠多少个
    heads=16, #多头
    mlp_dim=2048,
    dropout=0.1,
    emb_dropout=0.1
)

# 创建一个形状为(1, 3, 224, 224)的随机图像张量
img = torch.randn(2, 3, 224, 224)

# 使用ViT模型进行前向传播，得到分类预测
preds = v(img)  # 形状为(1, 1000)的分类分数，过一个softmax就是分类预测
print(preds.shape)