YOLO视觉检测应用 个性化模型篇

YOLO个性化模型

在YOLO的应用中，有很多时候对模型需要个性化的修改，以便得到更好的结果。
比如CBAM和GAM模块。
本文以这两个模块为例，介绍下YOLO个性化模型的应用。

Note:
本文主要关注两个模块的应用。
关于两个模型的具体情况，请查看相应的论文。

CBAM模块

CBAM全称 Convolutional Block Attention Module (卷积块注意力模块)。

CBAM论文

https://arxiv.org/abs/1807.06521

CBAM官方Github

https://github.com/Jongchan/attention-module

CBAM总体视图

其由Channel Attention Module(通道注意力模块)和Spatial Attenion Module(空间注意力模块)两个模块组成。

• Channel Attention Module(通道注意力模块):用于显式地建模通道之间的依赖关系
  
• Spatial Attenion(空间注意力)：用于显式地建模空间位置之间的依赖关系
  

由于CBAM是一个轻量级的通用模块，因此可以无缝地集成到任何CNN的架构中。可用性非常广泛。

Note:
具体的请参见CBAM论文

GAM模块

GAM(Global Attention Mechanism)全局注意力机制和CBAM模块还是比较类似的，同样使用了通道注意力和空间注意力，不同的是处理方式不一样。

GAM论文

http://arxiv.org/abs/2112.05561v1

GAM官方Github

https://github.com/Jongchan/attention-module

代码

为应用个性化的模型，我们首先准备代码
复制下列代码，存储为CBAM.py(也可以在网上找你觉得合适的代码，github上有很多)

######################  CBAM  GAM  ####     START   by  AI&CV  ###############################

import torch
from torch import nn
from torch.nn import init
import torch.nn.functional as F

from ultralytics.nn.modules.conv import Conv


class ChannelAttention(nn.Module):
    # Channel-attention module https://github.com/open-mmlab/mmdetection/tree/v3.0.0rc1/configs/rtmdet
    def __init__(self, channels: int) -> None:
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
        self.act = nn.Sigmoid()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * self.act(self.fc(self.pool(x)))


class SpatialAttention(nn.Module):
    # Spatial-attention module
    def __init__(self, kernel_size=7):
        super().__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.cv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.act = nn.Sigmoid()

    def forward(self, x):
        return x * self.act(self.cv1(torch.cat([torch.mean(x, 1, keepdim=True), torch.max(x, 1, keepdim=True)[0]], 1)))


class CBAM(nn.Module):
    # Convolutional Block Attention Module
    def __init__(self, c1, c2, kernel_size=7):  # ch_in, kernels
        super().__init__()
        self.channel_attention = ChannelAttention(c2)
        self.spatial_attention = SpatialAttention(kernel_size)

    def forward(self, x):
        return self.spatial_attention(self.channel_attention(x))




def channel_shuffle(x, groups=2):  ##shuffle channel
    # RESHAPE----->transpose------->Flatten
    B, C, H, W = x.size()
    out = x.view(B, groups, C // groups, H, W).permute(0, 2, 1, 3, 4).contiguous()
    out = out.view(B, C, H, W)
    return out


class GAM_Attention(nn.Module):
    def __init__(self, c1, c2, group=True, rate=4):
        super(GAM_Attention, self).__init__()

        self.channel_attention = nn.Sequential(
            nn.Linear(c1, int(c1 / rate)),
            nn.ReLU(inplace=True),
            nn.Linear(int(c1 / rate), c1)
        )

        self.spatial_attention = nn.Sequential(

            nn.Conv2d(c1, c1 // rate, kernel_size=7, padding=3, groups=rate) if group else nn.Conv2d(c1, int(c1 / rate),
                                                                                                     kernel_size=7,
                                                                                                     padding=3),
            nn.BatchNorm2d(int(c1 / rate)),
            nn.ReLU(inplace=True),
            nn.Conv2d(c1 // rate, c2, kernel_size=7, padding=3, groups=rate) if group else nn.Conv2d(int(c1 / rate), c2,
                                                                                                     kernel_size=7,
                                                                                                     padding=3),
            nn.BatchNorm2d(c2)
        )

    def forward(self, x):
        b, c, h, w = x.shape
        x_permute = x.permute(0, 2, 3, 1).view(b, -1, c)
        x_att_permute = self.channel_attention(x_permute).view(b, h, w, c)
        x_channel_att = x_att_permute.permute(0, 3, 1, 2)
        # x_channel_att=channel_shuffle(x_channel_att,4) #last shuffle
        x = x * x_channel_att

        x_spatial_att = self.spatial_attention(x).sigmoid()
        x_spatial_att = channel_shuffle(x_spatial_att, 4)  # last shuffle
        out = x * x_spatial_att
        # out=channel_shuffle(out,4) #last shuffle
        return out


######################  CBAM  GAM  ####     end   by  AI&CV  ###############################

建立个性化模型

源代码修改

建议先备份修改的文件/文件夹以备不时之需

在\Ultralytics\nn下建立ExtraModules文件夹。
将前面的cbam.py放进去,同时建立__init__.py,内容如下

from .cbam import *

修改task.py文件,加入

from ultralytics.nn.ExtraModules import *

然后找到程序def parse_model(d, ch, verbose=True):,在其中的

base_modules = frozenset(
        {
        }
)

中加入CBAM,GAM_Attention(名称与之前的cbam.py中class一致)

yaml文件准备

将模块插入你希望的地方，比如
GAM:

# YOLO11n head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 6], 1, Concat, [1]] # cat backbone P4
  - [-1, 2, C3k2, [512, False]] # 13

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)
  - [-1, 1, GAM_Attention, [256]] # 17

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 13], 1, Concat, [1]] # cat head P4
  - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium)
  - [-1, 1, GAM_Attention, [512]] # 21

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 10], 1, Concat, [1]] # cat head P5
  - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large)
  - [-1, 1, GAM_Attention, [1024]] # 25

  - [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)

CBAM:

# YOLO11n head
head:
  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 6], 1, Concat, [1]] # cat backbone P4
  - [-1, 2, C3k2, [512, False]] # 13

  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
  - [-1, 2, C3k2, [256, False]] # 16 (P3/8-small)

  - [-1, 1, Conv, [256, 3, 2]]
  - [[-1, 13], 1, Concat, [1]] # cat head P4
  - [-1, 2, C3k2, [512, False]] # 19 (P4/16-medium)

  - [-1, 1, Conv, [512, 3, 2]]
  - [[-1, 10], 1, Concat, [1]] # cat head P5
  - [-1, 2, C3k2, [1024, True]] # 22 (P5/32-large)
  
  - [16, 1, CBAM, [256]] # 23
  - [19, 1, CBAM, [512]] # 24
  - [22, 1, CBAM, [1024]] # 25

  - [[16, 19, 22], 1, Detect, [nc]] # Detect(P3, P4, P5)

Note:
以上仅为示例，请按照实际情况使用

运行个性化模型

之前在YOLO视觉检测应用入门篇（四）中，提供的代码中已经存在使用个性化模型的部分，只是先注释掉了。

# cbam_model= r"path\to\yolo11_CBAM.yaml" # 个性化的模型   

# model = YOLO(cfg.cbam_model).load(cfg.pretrained_weights)
# model = YOLO(cfg.cbam_model)

完整的代码请参阅：YOLO视觉检测应用入门篇（四）

你也可使用我在Github上的开源项目YOLO Train GUI，同样支持个性化模型的应用。
YOLO训练控制台

https://github.com/abslut/yolotraingui/releases

运行确认

GAM：
运行时可见GAM模块已经在预定的地方使用。

CBAM：
运行时可见CBAM模块已经在预定的地方使用。

小结

经过上面的操作，相信已经可以将个性化的YOLO模型运用起来了。
但是具体的效果，可能还需要不断的摸索。
㊗️好运！