ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

transformer用于图像分类

transformer用于图像分类 计算机视觉与transformerVision transformer1、VIT:1原理将图像分割成小块通过线性变换得到patch embedding加上位置编码输入到encoder中最后用分类头进行分类。2代码部分patch embeddingposition encodingtransformer encoder blockMuti-head attentionMLP最终分类头代码步骤patch embeddingposition encodingtransformer encoder blockMuti-head attentionMLP最终分类头代码步骤import torchfrom torch import nnfrom einops import rearrange#用于张量操作第一部分转变为编码class PatchEmbedding(nn.Module):def __init__(self, img_size32, patch_size4, in_channels3, embed_dim128):将图像分割为小块并线性投影到嵌入空间参数img_size: 输入图像尺寸假设为正方形patch_size: 每个小块的尺寸in_channels: 输入通道数RGB为3embed_dim: 嵌入维度embed_dim 全称 embedding dimension嵌入维度表示每个图像块patch被映射到的向量维度模型内部特征表示的统一维度所有 Transformer 层的输入/输出维度super().__init__()self.img_size img_sizeself.patch_size patch_size#将参数赋予给变量# 计算分块数量32/48 → 8x864个patchself.num_patches (img_size // patch_size) ** 2平方就得到了数量# 定义卷积层代替线性投影更高效self.projection nn.Conv2d(in_channelsin_channels,out_channelsembed_dim,kernel_sizepatch_size,stridepatch_size)这里不是很理解为什么要这么做呢为什么卷积核和步长要这样设置def forward(self, x):#定义前向传播将上边的函数用起来# x形状: (batch_size, 3, 32, 32)x self.projection(x)# 投影后形状: (batch_size, 128, 8, 8)这里的128就是转换后的特征维度x rearrange(x, b c h w - b (h w) c)# 展平为序列形状变为 (batch_size, 64, 128)一直到这里一张图由原来的3*32*32变为了64*128也就是真正的变成了64个向量。return xbatch_size, 64, 128第二部分多头注意力机制class MultiHeadAttention(nn.Module):def __init__(self, embed_dim128, num_heads4):super().__init__()self.num_heads num_headsself.head_dim embed_dim // num_heads#计算每个注意力头的维度128/432每个注意力头处理32维信息# 合并计算QKV的线性投影self.qkv nn.Linear(embed_dim, embed_dim * 3)#使用一个线性层输入维度是128输出是128*3也就是一个线性层同时生成Q/K/V然后将输出分割为三部分这样可以减少计算self.attention_dropout nn.Dropout(0.1)#注意力分数计算完成后使用随机关闭一些权重。self.projection nn.Linear(embed_dim, embed_dim)#将多头注意力后的结果进行投影保持维度一致性前期工作完成开始使用函数def forward(self, x):batch_size, seq_len, embed_dim x.shape#前面x的形状是bactch*64*128# 生成QKV形状 (batch_size, seq_len, 3*embed_dim)qkv self.qkv(x)在这里变成了batch*64*3*128# 拆分为多头形状 (3, batch_size, num_heads, seq_len, head_dim)将原先的batch*64*3*128变成batch*64*3*4*32的拆成3*batch*64*4*32qkv rearrange(qkv, b s (n h d) - n b h s d, n3, hself.num_heads)#输出为3*batch*4*64*32q, k, v qkv[0], qkv[1], qkv[2]#现在对三个矩阵取位置得到Q/K/V对应的矩阵形式# 计算注意力分数scores torch.matmul(q, k.transpose(-1, -2)) / (self.head_dim ** 0.5)#这里是在计算Q与K的点积相似度。Q的维度batch*4*64*32K的维度也是batch*4*64*32transpose(-1, -2)将K的最后两个维度可以计算点积Scores的形状是batch*4*64*64attention torch.softmax(scores, dim-1)#将scores转变为概率分布每个问询变量与key的注意力权重和为1dim-1表示在最后一个维度上进行softmax也就是在64上进行主要是确保每个问询变量与key的注意力权重和为1。得到的应该是batch*4*64*64个。attention self.attention_dropout(attention)#这里是为了防止过拟合# 加权求和x torch.matmul(attention, v)#attention的形状是batch*4*64*64v的形状是batch*4*64*32进行点积计算也就是用注意力权重对v进行加权平均得到新的V值x rearrange(x, b h s d - b s (h d))# 合并多头这里X的形状是batch*4*64*32转变成为batch*64*128也就是将多头的信息合并了x self.projection(x)#这里将上边合并后的信息进行线性投影将多头的信息融合return xbatch*64*128第三部分定义前向传播class MLP(nn.Module):#开始定义transformer中的前向传播def __init__(self, embed_dim128, hidden_dim512):super().__init__()self.net nn.Sequential(nn.Linear(embed_dim, hidden_dim),nn.GELU(),# ViT中常用GELU激活函数nn.Dropout(0.1),nn.Linear(hidden_dim, embed_dim),nn.Dropout(0.1))这里定义了一个线性层函数def forward(self, x):return self.net(x)#运用前边的函数进行前向传播batch*64*128第四部分定义BLOCKclass TransformerBlock(nn.Module):#开始定义transformer的编码器块def __init__(self, embed_dim128, num_heads4):super().__init__()self.norm1 nn.LayerNorm(embed_dim)#这里应该是进行了归一化吧self.attn MultiHeadAttention(embed_dim, num_heads)#这里用的是前边的类self.norm2 nn.LayerNorm(embed_dim)self.mlp MLP(embed_dim)#这里用到的是前面的前向传播网络def forward(self, x):# 残差连接 层归一化Pre-Norm结构x x self.attn(self.norm1(x))x x self.mlp(self.norm2(x))return xbatch*64*128第五部分开始使用前面的东西class VisionTransformer(nn.Module):def __init__(self, num_classes10, depth4):super().__init__()self.patch_embed PatchEmbedding()#将图片进行编码batch*64*128embed_dim 128# 可学习的分类token[CLS] token它的作用是汇总整个序列的信息用于最终的分类任务self.cls_token nn.Parameter(torch.randn(1, 1, embed_dim))生成一个形状为1*1*128的张量1表示初始批次大小为11表示序列中只有一个可学习的tokenembed_dim表示与图像的编码维度一致也就是使用正态分布进行标准化。nn.Parameter表示该张量是模型的可训练参数。# 位置编码可学习参数self.pos_embed nn.Parametertorch.randn(1, self.patch_embed.num_patches 1, embed_dim))#这里为什么是65是因为还要为前边的cls_token进行编码这里的1表示所有样本共享同一组位置编码不随批次变化而变化# 创建Transformer编码器堆叠self.blocks nn.Sequential(*[TransformerBlock() for _ in range(depth)])#depth4nn.Sequential(*[block1, block2, block3])等价于nn.Sequential(block1, block2, block3)for _ in range(depth)重复创建 depth 个相同的 TransformerBlock。batch*64*128# 分类头self.head nn.Sequential(nn.LayerNorm(embed_dim),nn.Linear(embed_dim, num_classes))输出为batch*10这里的输入的形状是batch*128定义前向传播开始使用前边的类def forward(self, x):batch_size x.shape[0]# 生成patch嵌入x self.patch_embed(x) #这个类下边定义的但用了前边的类形状 (batch_size, 64, 128)# 添加分类tokencls_token self.cls_token.expand(batch_size, -1, -1)#将cls_token的形状变为实际批次的大小x torch.cat([cls_token, x], dim1) # 形状 (batch_size, 65, 128)将cls_token拼接到x中去。# 添加位置编码x self.pos_embed#每个x都在自身编码的基础上加上位置编码# 通过Transformer编码器x self.blocks(x)# 取出分类token的特征只有这个是可以用来回归或者使用的特征cls_token_final x[:, 0]这个意思就是说在(batch_size, 65, 128中取出也就是每65个取出第一个张量得到batch_size, 128# 分类头return self.head(cls_token_final)得到最终的分类结果# 使用示例if __name__ __main__:# 创建虚拟输入batch_size4dummy_img torch.randn(4, 3, 32, 32) # 4张32x32的RGB图片# 初始化ViT模型vit VisionTransformer(num_classes10)# 前向传播output vit(dummy_img)print(输出形状:, output.shape) # 应该得到 torch.Size([4, 10]
返回列表