YOLOV8改进：如何增加注意力模块？（以CBAM模块为例）

2023-12-25 14:25| 来源: 网络整理| 查看: 265

YOLOV8改进：如何增加注意力模块？（以CBAM模块为例）前言YOLOV8nn文件夹modules.pytask.py models文件夹总结

前言

因为毕设用到了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':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:

【本文地址】

公司简介

联系我们