基于Swin Transformer与CBAM增强的YOLO混合感知模型设计与实战

1次阅读

没有评论

共计 2045 个字符，预计需要花费 6 分钟才能阅读完成。

传统 CNN-based 目标检测模型（如 YOLO 系列）在复杂场景下面临两个核心问题：

长距离依赖建模不足：卷积操作的局部感受野特性导致模型难以捕捉全局上下文关系。在 COCO 数据集的验证集上，YOLOv5s 对跨图像区域的交互物体（如手持物品与人）的检测 AP 仅达到 23.1%，比人类标注差距达 16.7 个百分点。
复杂背景干扰：当目标与背景颜色 / 纹理相似时，传统模型容易产生误检。可视化热图显示，YOLOv5 在树叶遮挡场景下的误检率高达 34%。

窗口注意力 ：将图像划分为 4×4 局部窗口（默认窗口大小 7×7），计算复杂度从标准 ViT 的 O(n²) 降至 O(n)。在 512×512 输入下，Swin- T 的计算量仅为 ViT- B 的 18%。
层级设计：通过 Patch Merging 实现 4 阶段特征图下采样（224×224→56×56→28×28→14×14），完美适配 YOLO 的多尺度检测需求。

通道注意力：使用 GAP+MLP 学习通道权重，公式：
$$M_c(F) = \sigma(MLP(AvgPool(F)) + MLP(MaxPool(F)))$$
空间注意力：卷积层捕捉空间关系，公式：
$$M_s(F) = \sigma(f^{7×7}([AvgPool(F); MaxPool(F)]))$$

实验表明，在 COCO 验证集上单独使用 CBAM 可使小目标 AP 提升 2.3%。

Swin-CBAM 骨干网络：
Stage1: 48-dim, 输出 160×160×48
Stage2: 96-dim, 输出 80×80×96
Stage3: 192-dim, 输出 40×40×192（接入 CBAM）
Stage4: 384-dim, 输出 20×20×384
YOLO 检测头适配：
采用 PAFPN 结构融合多尺度特征
使用 1×1 卷积统一通道数到 256

在 Stage3 输出后执行特征重整：
$$F_{out} = SwinBlock(CBAM(F_{in})) + F_{in}$$

class SwinCBAMBlock(nn.Module):
    def __init__(self, dim, window_size=7):
        super().__init__()
        self.swin = SwinTransformerBlock(dim, window_size)
        self.cbam = CBAM(gate_channels=dim)

    def forward(self, x):
        # 输入尺寸: (B, C, H, W)
        x = x.permute(0, 2, 3, 1)  # 转为(B,H,W,C)
        x = self.swin(x)
        x = x.permute(0, 3, 1, 2)  # 转回(B,C,H,W)
        return self.cbam(x)

class YOLOHeadAdapter(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.convs = nn.ModuleDict({'p3': nn.Conv2d(in_channels[0], 256, 1),
            'p4': nn.Conv2d(in_channels[1], 256, 1),
            'p5': nn.Conv2d(in_channels[2], 256, 1)
        })

    def forward(self, features):
        return [self.convs['p3'](features[0]), 
                self.convs['p4'](features[1]),
                self.convs['p5'](features[2])]

模型	AP@0.5	AP@0.5:0.95	参数量(M)	FPS
YOLOv5s	37.2	20.4	7.2	142
本方案(Swin-T)	41.6(+4.4)	23.7(+3.3)	8.1	119

YOLOv5s: 1.8GB
本方案: 2.3GB（增加 27.8%）

多尺度训练：
初始 lr=0.001，每增加一个尺度乘以 0.8 系数
warmup 阶段保持单尺度训练 3epoch
混合精度训练：
梯度裁剪阈值设为 3.0（默认 5.0 易导致 NaN）
使用 AdamW 优化器时设置 eps=1e-4

ONNX 导出：

torch.onnx.export(model, 
                dummy_input,
                "model.onnx",
                dynamic_axes={"images": {0: "batch"}, 
                             "outputs": {0: "batch"}})