您好,欢迎访问欧亿体育!

020-88886892

全国咨询热线

您现在所在位置: 主页 > 通知资讯

YOLOV8改进:如何增加注意力模块?(以CBAM模块为例)

更新时间:2024-01-21

YOLOV8改进:如何增加注意力模块?(以CBAM模块为例)

前言

因为毕设用到了YOLO,鉴于最近V8刚出,因此考虑将注意力机制加入到v8中。

YOLOV8

代码地址:YOLOV8官方代码

使用pip安装或者clone到本地,在此不多赘述了。下面以使用pip安装ultralytics包为例介绍。
进入ultralytics文件夹

nn文件夹

再进入nn文件夹。

-- modules.py:在里面存放着各种常用的模块,如:Conv,DWConv,ConvTranspose,TransformerLayer,Bottleneck等
-- tasks.py: 在里面导入了modules中的基本模块组建model,根据不同的下游任务组建不同的model。
modules.py

在该文件中,我们可以写入自己的注意力模块,或者使用V8已经提供的CBAM模块(见代码的CBAM类)

"""
通道注意力模型: 通道维度不变,压缩空间维度。该模块关注输入图片中有意义的信息。
1)假设输入的数据大小是(b,c,w,h)
2)通过自适应平均池化使得输出的大小变为(b,c,1,1)
3)通过2d卷积和sigmod激活函数后,大小是(b,c,1,1)
4)将上一步输出的结果和输入的数据相乘,输出数据大小是(b,c,w,h)。
"""
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)))

"""
空间注意力模块:空间维度不变,压缩通道维度。该模块关注的是目标的位置信息。
1) 假设输入的数据x是(b,c,w,h),并进行两路处理。
2)其中一路在通道维度上进行求平均值,得到的大小是(b,1,w,h);另外一路也在通道维度上进行求最大值,得到的大小是(b,1,w,h)。
3) 然后对上述步骤的两路输出进行连接,输出的大小是(b,2,w,h)
4)经过一个二维卷积网络,把输出通道变为1,输出大小是(b,1,w,h)
4)将上一步输出的结果和输入的数据x相乘,最终输出数据大小是(b,c,w,h)。
"""
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, kernel_size=7):  # ch_in, kernels
        super().__init__()
        self.channel_attention = ChannelAttention(c1)
        self.spatial_attention = SpatialAttention(kernel_size)
        
        

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

如果使用V8的CBAM模块,则不需要更改modules.py的内容。如果使用自己的注意力模块,只需要在该文件后面添加对应的代码即可。

task.py

在该文件中,通过import modules.py文件中的模块来构建模型。
在文件开头导入需要的模块,可以看到modules中的很多模块在v8中并没有用到。我们在最后添加对应的CBAM模块。

from ultralytics.nn.modules import (C1, C2, C3, C3TR, SPP, SPPF, Bottleneck, BottleneckCSP, C2f, C3Ghost, C3x, Classify,
                                    Concat, Conv, ConvTranspose, Detect, DWConv, DWConvTranspose2d, Ensemble, Focus,
                                    GhostBottleneck, GhostConv, Segment, CBAM)

之后修改对应的parse_model方法(对应428行)
添加分支elif m is CBAM:,具体代码如下:

def parse_model(d, ch, verbose=True):  # model_dict, input_channels(3)
    # Parse a YOLO model.yaml dictionary
    if verbose:
        LOGGER.info(f"\n{'':>3}{'from':>20}{'n':>3}{'params':>10}  {'module':<45}{'arguments':<30}")
    nc, gd, gw, act = d['nc'], d['depth_multiple'], d['width_multiple'], d.get('activation')
    if act:
        Conv.default_act = eval(act)  # redefine default activation, i.e. Conv.default_act = nn.SiLU()
        if verbose:
            LOGGER.info(f"{colorstr('activation:')} {act}")  # print
    ch = [ch]
    layers, save, c2 = [], [], ch[-1]  # layers, savelist, ch out
    for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']):  # from, number, module, args
        m = eval(m) if isinstance(m, str) else m  # eval strings
        for j, a in enumerate(args):
            # TODO: re-implement with eval() removal if possible
            # args[j] = (locals()[a] if a in locals() else ast.literal_eval(a)) if isinstance(a, str) else a
            with contextlib.suppress(NameError):
                args[j] = eval(a) if isinstance(a, str) else a  # eval strings

        n = n_ = max(round(n * gd), 1) if n > 1 else n  # depth gain
        if m in (Classify, Conv, ConvTranspose, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, Focus,
                 BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x):
            c1, c2 = ch[f], args[0]
            if c2 != nc:  # if c2 not equal to number of classes (i.e. for Classify() output)
                c2 = make_divisible(c2 * gw, 8)

            args = [c1, c2, *args[1:]]
            if m in (BottleneckCSP, C1, C2, C2f, C3, C3TR, C3Ghost, C3x):
                args.insert(2, n)  # number of repeats
                n = 1
        elif m is nn.BatchNorm2d:
            args = [ch[f]]
        elif m is Concat:
            c2 = sum(ch[x] for x in f)
        elif m in (Detect, Segment):
            args.append([ch[x] for x in f])
            if m is Segment:
                args[2] = make_divisible(args[2] * gw, 8)
        elif m is CBAM:
            """
            ch[f]:上一层的
            args[0]:第0个参数
            c1:输入通道数
            c2:输出通道数
            """
            c1, c2 = ch[f], args[0]
            # print("ch[f]:",ch[f])
            # print("args[0]:",args[0])
            # print("args:",args)
            # print("c1:",c1)
            # print("c2:",c2)
            if c2 != nc:  # if c2 not equal to number of classes (i.e. for Classify() output)
                c2 = make_divisible(c2 * gw, 8)
            args = [c1,*args[1:]]
        else:
            c2 = ch[f]
        
        m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args)  # module
        t = str(m)[8:-2].replace('__main__.', '')  # module type
        m.np = sum(x.numel() for x in m_.parameters())  # number params
        m_.i, m_.f, m_.type = i, f, t  # attach index, 'from' index, type
        if verbose:
            LOGGER.info(f'{i:>3}{str(f):>20}{n_:>3}{m.np:10.0f}  {t:<45}{str(args):<30}')  # print
        save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1)  # append to savelist
        layers.append(m_)
        if i == 0:
            ch = []
        ch.append(c2)
    return nn.Sequential(*layers), sorted(save)

注意传入的参数为上一层输出,要注意CBAM模块的参数和传入参数的对应。读者可以自行print比较。

models文件夹

返回上一级目录,进入models文件夹。
可以看到该文件夹中还有v5、v3对应的模型配置文件,所以也可以使用该包进行v5和v3的训练。
进入v8文件夹

打开对应的yolov8.yaml,如下所示。该文件是V8对应的配置文件,里面包括了类别数,模型大小(n,s,m,l,x),backbone和head。

# Ultralytics YOLO  

联系电话

020-88886892

返回顶部