Conformer代码讲解

主要讲重点&＃xff1a;CNN与transformer模块是怎么融合的

输入图片卷积后为&＃xff08;b,256,56,56&＃xff09;。四倍下采样后为&＃xff08;b&＃xff0c;384&＃xff0c;14&＃xff0c;14&＃xff09;&＃xff0c;再加上&＃xff08;1&＃xff0c;384&＃xff09;维度的class_token成为&＃xff08;b&＃xff0c; 197&＃xff0c;384&＃xff09;

for i in range(2, self.fin_stage):x, x_t &＃61; eval(&＃39;self.conv_trans_&＃39; &＃43; str(i))(x, x_t)
# 这里要重复2-12次&＃xff0c;输入是特征图x&＃xff08;b,256,56,56&＃xff09;与token x_t(b,197,384)&＃xff0c;输出也是这两个部分。下面拆开讲解循环部分&＃xff1a;

2.初始下采样

def forward(self, x, x_t):x, x2 &＃61; self.cnn_block(x) # 第一次后维度x(b,256,56,56) x2(b,64,56,56) i&＃61;5时为 (512,28,28)(128, 28, 28)# self.cnn_block作用是下采样&＃xff0c;在i循环中&＃xff08;2-12&＃xff09;&＃xff0c;2-4不变&＃xff0c;5-8不变&＃xff0c;9不变&＃xff0c;12-12不变

3.CNN–>Trans

_, _, H, W &＃61; x2.shapex_st &＃61; self.squeeze_block(x2, x_t) x_t &＃61; self.trans_block(x_st &＃43; x_t)# 特征图x2经过2次卷积-->(b&＃xff0c;196,384)&＃xff0c;叠加x_t的第一维&＃xff0c;成为 x_st(b, 197, 384)。x_st 与 x_t相加&＃xff0c;输入Trans模块&＃xff0c;得到 x_t维度不变。


self.squeeze_block&＃xff1a;(conv_project): Conv2d(128, 384, kernel_size&＃61;(1, 1), stride&＃61;(1, 1))(sample_pooling): AvgPool2d(kernel_size&＃61;2, stride&＃61;2, padding&＃61;0)(ln): LayerNorm((384,), eps&＃61;1e-06, elementwise_affine&＃61;True)(act): GELU()def forward(self, x, x_t):x &＃61; self.conv_project(x) # [N, C, H, W]x &＃61; self.sample_pooling(x).flatten(2).transpose(1, 2)x &＃61; self.ln(x)x &＃61; self.act(x)x &＃61; torch.cat([x_t[:, 0][:, None, :], x], dim&＃61;1)return x

4.Trans–>CNN

x_t_r &＃61; self.expand_block(x_t, H // self.dw_stride, W // self.dw_stride) x &＃61; self.fusion_block(x, x_t_r, return_x_2&＃61;False)
# 将token embed 进行双线性插值&＃xff0c;得到 x_t_r(b&＃xff0c; 64, 56, 56])&＃xff0c;增大了分辨率
# 变成矩阵的 x_t_r 再加回到特征图x&＃xff0c;得到xreturn x, x_t


self.expand_block&＃xff1a; (conv_project): Conv2d(384, 128, kernel_size&＃61;(1, 1), stride&＃61;(1, 1))(bn): BatchNorm2d(128, eps&＃61;1e-06, momentum&＃61;0.1, affine&＃61;True, track_running_stats&＃61;True)(act): ReLU()def forward(self, x, H, W):B, _, C &＃61; x.shape# [N, 197, 384] -> [N, 196, 384] -> [N, 384, 196] -> [N, 384, 14, 14]x_r &＃61; x[:, 1:].transpose(1, 2).reshape(B, C, H, W)x_r &＃61; self.act(self.bn(self.conv_project(x_r)))return F.interpolate(x_r, size&＃61;(H * self.up_stride, W * self.up_stride))

self.fusion_block&＃xff1a;(conv1): Conv2d(512, 128, kernel_size&＃61;(1, 1), stride&＃61;(1, 1), bias&＃61;False)(bn1): BatchNorm2d(128, eps&＃61;1e-06, momentum&＃61;0.1, affine&＃61;True, track_running_stats&＃61;True)(act1): ReLU(inplace&＃61;True)(conv2): Conv2d(128, 128, kernel_size&＃61;(3, 3), stride&＃61;(1, 1), padding&＃61;(1, 1), bias&＃61;False)(bn2): BatchNorm2d(128, eps&＃61;1e-06, momentum&＃61;0.1, affine&＃61;True, track_running_stats&＃61;True)(act2): ReLU(inplace&＃61;True)(conv3): Conv2d(128, 512, kernel_size&＃61;(1, 1), stride&＃61;(1, 1), bias&＃61;False)(bn3): BatchNorm2d(512, eps&＃61;1e-06, momentum&＃61;0.1, affine&＃61;True, track_running_stats&＃61;True)(act3): ReLU(inplace&＃61;Truedef forward(self, x, x_t&＃61;None, return_x_2&＃61;True):residual &＃61; x # (b,256, 56, 56)x &＃61; self.conv1(x)x &＃61; self.bn1(x)if self.drop_block is not None:x &＃61; self.drop_block(x)x &＃61; self.act1(x)self.conv2(x &＃43; x_t)x &＃61; self.conv2(x) if x_t is not None else x &＃61; self.conv2(x)x &＃61; self.bn2(x)if self.drop_block is not None:x &＃61; self.drop_block(x)x2 &＃61; self.act2(x)x &＃61; self.conv3(x2)x &＃61; self.bn3(x)if self.drop_block is not None:x &＃61; self.drop_block(x)if self.drop_path is not None:x &＃61; self.drop_path(x)if self.res_conv:residual &＃61; self.residual_conv(residual)residual &＃61; self.residual_bn(residual)x &＃43;&＃61; residualx &＃61; self.act3(x)return x

5. 分类阶段

在i&＃61;2-12时&＃xff0c;输出维度依次为:
(1) x :(b,256,56,56) (b,512,28,28) (b,1024,14,14) (b,1024,7,7)
(2)x_t: (b,197,384) …(b,197,384)
x:(b,1024,7,7) —avgPool–>(b,1024)----conv_cls_head—>(b,1000)
x_t: (b,197,384)—取第一维—>&＃xff08;b,1,384&＃xff09;—trans_cls_head–>(b,1000)
最后结果取两个的平均