1. PVT_V1的设计动机与核心创新
当你第一次看到Vision Transformer(ViT)时,可能会被它处理图像的方式惊艳到——把图像切成小块当作序列处理。但实际用起来就会发现,ViT在密集预测任务(比如目标检测、语义分割)中表现平平。这就像给你一把瑞士军刀,却发现它切牛排不如专业牛排刀顺手。
PVT_V1的诞生正是为了解决ViT的两个关键痛点。首先是单尺度特征图问题。想象你要装修房子,ViT只给你提供了一种比例的设计图纸,而传统CNN(比如ResNet)却能提供从整体布局到插座位置的各级详图。PVT_V1通过金字塔结构,让Transformer也能输出类似CNN的多级特征图。
更棘手的是计算效率问题。处理一张800px的图片时,ViT需要计算全部1600个patch(假设patch大小为20x20)之间的注意力关系,这会产生256万次计算!PVT_V1的解决方案相当巧妙——用空间缩减注意力(SRA)机制把计算量压缩到原来的1/64,就像用缩略图快速找出重点区域,再对原图精细处理。
2. 网络架构全景解读
2.1 金字塔结构设计
PVT_V1的整体架构很容易让人联想到ResNet,这种刻意对齐的设计让替换现有模型变得轻松。来看具体的数据流动过程:
- Stage 1:输入224x224图像 → 4x4卷积(stride=4) → 56x56特征图
- Stage 2:56x56输入 → 3x3卷积(stride=2) → 28x28特征图
- Stage 3:28x28 → 3x3卷积(stride=2) → 14x14特征图
- Stage 4:14x14 → 3x3卷积(stride=2) → 7x7特征图
每个stage的通道数也在递增,典型配置是[64, 128, 320, 512]。这种设计让下游任务可以像使用ResNet那样,自由组合不同层级的特征。
2.2 关键组件拆解
每个stage的核心是若干个Transformer Block,其结构比ViT多了一个重要部件:
class Block(nn.Module): def __init__(self, dim, num_heads, sr_ratio=1, ...): super().__init__() self.norm1 = nn.LayerNorm(dim) self.attn = Attention(dim, num_heads, sr_ratio) # 关键改动在这里 self.norm2 = nn.LayerNorm(dim) self.mlp = Mlp(dim) def forward(self, x, H, W): x = x + self.attn(self.norm1(x), H, W) # 带空间信息的注意力 x = x + self.mlp(self.norm2(x)) # 标准MLP return x与ViT最大的区别在于Attention模块需要接收特征图的宽高信息(H,W),这是实现空间缩减的关键。下面我们就深入这个最核心的创新点。
3. 空间缩减注意力(SRA)实现详解
3.1 原版注意力的问题
标准Transformer的注意力计算复杂度是O(N²),其中N是patch数量。对于56x56的特征图,N=3136,计算量达到惊人的:
3136 × 3136 ≈ 980万次计算这还只是单个注意力头在单个样本上的计算量!PVT_V1通过三步实现计算优化:
- 空间缩减:用卷积压缩特征图尺寸
- 键值生成:在低分辨率特征上生成K、V
- 查询保持:仍在原始分辨率上生成Q
3.2 代码逐行解析
来看Attention类的关键实现(以sr_ratio=8为例):
def forward(self, x, H, W): B, N, C = x.shape # 输入形状 (1, 3136, 64) # 生成Q向量(保持原始分辨率) q = self.q(x).reshape(B, N, self.num_heads, C//self.num_heads) q = q.permute(0, 2, 1, 3) # (1, 1, 3136, 64) # 空间缩减关键步骤 x_ = x.permute(0, 2, 1).reshape(B, C, H, W) # 转图像格式 (1,64,56,56) x_ = self.sr(x_) # 用8x8卷积压缩 (1,64,7,7) x_ = x_.reshape(B, C, -1).permute(0, 2, 1) # (1,49,64) x_ = self.norm(x_) # 生成K、V向量 kv = self.kv(x_).reshape(B, -1, 2, self.num_heads, C//self.num_heads) kv = kv.permute(2, 0, 3, 1, 4) # (2,1,1,49,64) k, v = kv[0], kv[1] # 各(1,1,49,64) # 注意力计算 attn = (q @ k.transpose(-2,-1)) * self.scale # (1,1,3136,49) attn = attn.softmax(dim=-1) x = (attn @ v).transpose(1,2).reshape(B,N,C) # (1,3136,64) return x计算量从980万次降到了约15万次(3136×49),效果提升约64倍!这种设计既保留了全局感知能力,又大幅降低了计算成本。
4. 特征变换全流程剖析
4.1 Patch Embedding实现细节
PVT_V1的patch嵌入比ViT更灵活,来看具体实现:
class PatchEmbed(nn.Module): def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=64): super().__init__() self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) self.norm = nn.LayerNorm(embed_dim) def forward(self, x): x = self.proj(x) # (1,3,224,224)->(1,64,56,56) x = x.flatten(2) # (1,64,3136) x = x.transpose(1, 2) # (1,3136,64) x = self.norm(x) return x, (56, 56) # 返回特征图尺寸有趣的是,后续stage的patch嵌入使用3x3卷积而非2x2,这样可以在下采样时更好地保留局部信息。例如Stage2的配置:
PatchEmbed(img_size=56, patch_size=3, stride=2, in_chans=64, embed_dim=128) # 56x56->28x284.2 位置编码的巧妙设计
PVT_V1的位置编码是可学习的参数,但有个特殊处理:
pos_embed = nn.Parameter(torch.zeros(1, 3136, 64)) # 可学习参数 # 在forward中处理不同输入尺寸 if H * W != self.patch_embed.num_patches: pos_embed = F.interpolate( pos_embed.reshape(1, 56, 56, -1).permute(0,3,1,2), size=(H,W), mode='bilinear' ).reshape(1,-1,H*W).permute(0,2,1)这种设计让模型可以处理可变尺寸输入,对目标检测等任务特别有用。我在实际使用中发现,相比ViT的固定位置编码,这种灵活设计使PVT_V1在迁移到不同分辨率时表现更稳定。
5. 完整模型实现与调参技巧
5.1 模型配置详解
PVT_V1提供多种预置配置,以pvt_small为例:
model = PyramidVisionTransformer( patch_size=4, embed_dims=[64, 128, 320, 512], # 各阶段通道数 num_heads=[1, 2, 5, 8], # 注意力头数 mlp_ratios=[8, 8, 4, 4], # MLP扩展系数 depths=[3, 4, 6, 3], # 各阶段block数 sr_ratios=[8, 4, 2, 1] # 空间缩减比率 )几个关键设计选择:
- 浅层用大sr_ratio:早期特征图尺寸大,更需要压缩
- 深层增加头数:高层语义需要更细粒度的注意力
- MLP比率递减:浅层需要更强的特征变换能力
5.2 实战训练技巧
基于在COCO数据集上的实测经验,分享几个调参要点:
学习率设置:
lr = 1e-4 * batch_size / 64 # 线性缩放规则权重衰减:
optimizer = AdamW(model.parameters(), lr=lr, weight_decay=0.05)数据增强:
transform = Compose([ RandomResizedCrop(224, scale=(0.2, 1.0)), RandomHorizontalFlip(), ColorJitter(0.4, 0.4, 0.4) ])
特别要注意的是,当迁移到下游任务时,建议先冻结stem和早期stage的参数,只微调高层block,这能有效防止过拟合。