Vision Transformer编码流程及代码详解前言传统CNN依靠卷积核局部滑动提取图像特征依赖归纳偏置局部性、平移不变性而Vision TransformerViT完全基于自注意力机制将图像拆分为序列Patch借用NLP Transformer架构完成全局特征建模。本文以工程最常用的ViT-B/16为例完整拆解从输入图像[3,224,224]到最终一维图像特征[1,768]的全流程维度变换、核心模块原理并附带可运行PyTorch完整实现代码逐行注释方便调试。一、ViT-B/16 参数含义B表示Base代表模型基础尺寸是轻量化常用版本编码器堆叠12层L(Large)大模型编码器24层H(Huge)超大模型编码器32层16代表Patch分块尺寸将224×224原图切分为16×16像素的小块常见Patch尺寸16、32、14ViT-L/14多用于高精度图像任务。二、前置基础核心模块维度运算基础理解ViT的关键是全程跟踪张量维度变化先回顾矩阵乘法规则向量a [1,768]× 权重矩阵W [768,512] 输出向量c [1,512]内维度必须相等输出维度由向量第一维、矩阵第二维决定。下面逐个介绍ViT编码器全部基础组件1. Linear 线性层公式yWxby Wx byWxb对输入特征做线性投影无非线性多用于Patch嵌入、多头注意力映射、分类头映射是维度升降维的核心层。2. GELU 激活函数高斯误差线性单元替代ReLU平滑非线性激活。解决ReLU梯度硬截断问题ViT前馈网络统一使用GELU提升特征拟合能力捕获图像复杂纹理、语义特征。3. Dropout 随机失活正则化手段训练阶段随机置零部分神经元输出推理阶段恢复完整权重。迫使模型不依赖局部少数神经元降低过拟合ViT在嵌入层、注意力输出、前馈层后均会添加Dropout。4. Layer Normalization 层归一化对单一样本自身特征维度做归一化区别于CNN常用的BatchNorm批次维度归一。Transformer系列标配稳定每层输入分布大幅加速深层模型收敛每层注意力、前馈网络前都会先做LN。5. Self-Attention 自注意力机制核心计算序列内每个Token与所有Token的相关性权重。输入序列中每个Patch Token互相计算相似度建模图像全局依赖远距离像素关联CNN很难做到。6. Multi-Head Attention 多头自注意力将特征通道均分N个注意力头每个头独立计算自注意力最后拼接所有头输出再线性融合。多个头并行捕捉不同维度、不同尺度的空间关联边缘、色块、全局轮廓单头注意力表达能力不足。7. FFN 前馈神经网络两层全连接GELU激活升维映射→激活→降维映射。独立作用于序列每一个Token对注意力输出特征做非线性特征变换增强模型表征能力。8. Residual Connection 残差连接OutputInputSubLayer(Input)Output Input SubLayer(Input)OutputInputSubLayer(Input)每层注意力、FFN外层包裹残差相加深层堆叠时避免梯度消失保证梯度跨层回流是12/24层深层ViT训练的基础。9. Positional Encoding 位置编码自注意力本身不感知序列顺序图像Patch打散后丢失空间位置信息。通过可学习位置编码ViT原生方案或正余弦编码生成和Patch Embedding同维度位置向量逐元素相加嵌入序列还原图像二维空间信息。三、ViT完整编码工作流程维度全程跟踪输入单张RGB图像张量[C3, H224, W224]batch_size1输入形状[1,3,224,224]步骤1图像切分Patch块按照Patch_size16切割原图横向分块224/1614224 / 16 14224/1614纵向同理14块总Patch数量14×1419614×1419614×14196每个Patch像素尺寸[3,16,16]全局图像拆分后得到196个独立图像小块。步骤2Patch Embedding 图像分块嵌入单个Patch[3,16,16]展平一维3×16×167683×16×167683×16×16768单Patch展平向量维度[768]全部196个Patch堆叠得到原始Patch序列[196, 768]批量维度扩展batch1张量形状[1, 196, 768]核心逻辑使用Conv2d卷积等价实现Patch切分线性投影工程上速度更快卷积核16步长16输出通道768卷积输出直接reshape为序列Token。步骤3拼接Class Token分类向量ViT新增一个专属分类TokenClass Token形状[1,1,768]拼接在Patch序列最前端原序列[1,196,768] Class Token → 新序列[1, 197, 768]模型最终全局图像特征从该Class Token提取对应文末输出[1,768]特征。步骤4叠加可学习位置编码创建与序列等长的位置编码参数[197,768]覆盖196个Patch 1个Class Token逐元素相加到嵌入序列注入空间位置信息。叠加后张量尺寸不变[1, 197, 768]随后经过Dropout做正则。步骤5堆叠12层Transformer EncoderViT-B核心编码层每层Encoder结构固定LN层 → 多头自注意力 残差连接 → LN层 → FFN前馈网络 残差连接逐层迭代计算全局注意力特征每层输入输出维度均保持[1,197,768]不变。单层Encoder数据流输入x [1,197,768]层归一化LN1 → 多头注意力MHA → x_attn x MHA(LN1(x)) 残差相加对x_attn做层归一化LN2 → FFN前馈网络 → x_out x_attn FFN(LN2(x_attn)) 残差相加x_out作为下一层Encoder输入12层循环结束最终输出编码后完整序列[1, 197, 768]步骤6提取全局图像特征目标输出[1,768]197个Token中第0位为Class Token代表整张图像聚合全局语义特征切片取出x[:, 0, :]形状[1, 768]即文章开头所说图像全局特征向量。后续分类任务可再接Linear层映射至类别数量检测/分割任务则取用全部Patch Token[:,1:,:]。四、ViT-B/16 完整PyTorch实现代码importtorchimporttorch.nnasnnimporttorch.nn.functionalasF# 超参数配置 ViT-B/16BATCH_SIZE1IMG_CHANNEL3IMG_SIZE224PATCH_SIZE16EMBED_DIM768# Base模型特征维度NUM_HEADS12# 多头注意力头数NUM_LAYERS12# Encoder层数MLP_HIDDEN3072# FFN隐藏层维度DROPOUT_RATE0.1# 1. 单层FFN前馈网络classFeedForward(nn.Module):def__init__(self):super().__init__()self.netnn.Sequential(nn.Linear(EMBED_DIM,MLP_HIDDEN),nn.GELU(),nn.Dropout(DROPOUT_RATE),nn.Linear(MLP_HIDDEN,EMBED_DIM),nn.Dropout(DROPOUT_RATE))defforward(self,x):returnself.net(x)# 2. 单层Transformer EncoderclassTransformerEncoderLayer(nn.Module):def__init__(self):super().__init__()self.norm1nn.LayerNorm(EMBED_DIM)self.attnnn.MultiheadAttention(EMBED_DIM,NUM_HEADS,dropoutDROPOUT_RATE,batch_firstTrue)self.norm2nn.LayerNorm(EMBED_DIM)self.ffnFeedForward()defforward(self,x):# 多头注意力 残差attn_out,_self.attn(queryself.norm1(x),keyself.norm1(x),valueself.norm1(x))xxattn_out# FFN 残差ffn_outself.ffn(self.norm2(x))xxffn_outreturnx# 3. 完整ViT-B/16 编码器classViT_B16_Encoder(nn.Module):def__init__(self):super().__init__()num_patches(IMG_SIZE//PATCH_SIZE)**2# 14*14196# Patch Embedding卷积替代分块线性投影self.patch_embednn.Conv2d(IMG_CHANNEL,EMBED_DIM,kernel_sizePATCH_SIZE,stridePATCH_SIZE)# 可学习Class Tokenself.cls_tokennn.Parameter(torch.randn(1,1,EMBED_DIM))# 可学习位置编码196patch 1cls_tokenself.pos_embednn.Parameter(torch.randn(1,num_patches1,EMBED_DIM))self.pos_dropnn.Dropout(DROPOUT_RATE)# 堆叠12层Encoderself.encoder_layersnn.Sequential(*[TransformerEncoderLayer()for_inrange(NUM_LAYERS)])self.norm_finalnn.LayerNorm(EMBED_DIM)defforward(self,img):# img输入 shape [B,3,224,224]Bimg.shape[0]# Step1 Patch Embedding [B,768,14,14] - [B,196,768]patch_featself.patch_embed(img)patch_featpatch_feat.flatten(2).transpose(1,2)# Step2 拼接Class Tokencls_tokensself.cls_token.expand(B,-1,-1)# [B,1,768]xtorch.cat([cls_tokens,patch_feat],dim1)# [B,197,768]# Step3 叠加位置编码Dropoutxxself.pos_embed xself.pos_drop(x)# Step4 12层Transformer编码xself.encoder_layers(x)xself.norm_final(x)# Step5 提取全局图像特征 cls_token [B,768]global_img_featx[:,0,:]returnglobal_img_feat# 测试流程if__name____main__:# 模拟输入图片 [1,3,224,224]test_imgtorch.randn(BATCH_SIZE,IMG_CHANNEL,IMG_SIZE,IMG_SIZE)modelViT_B16_Encoder()feat_outmodel(test_img)print(输入图像尺寸:,test_img.shape)print(输出全局图像特征尺寸:,feat_out.shape)# 输出结果torch.Size([1, 768])和文中结论完全对应