修改
在YOLOv8中ultralytics->nn->tasks.py文件,导入CBAM模块,在
from ultralytics.nn.modules import # 添加CBAM模块
添加
在
def parse_model(d, ch, verbose=True): # 中添加CBAM模块
添加 elif m is CBAM:
if m in {BottleneckCSP, C1, C2, C2f, C2fAttn, C3, C3TR, C3Ghost, C3x, RepC3}:
args.insert(2, n) # number of repeats
n = 1
elif m is AIFI:
args = [ch[f], *args]
elif m in {HGStem, HGBlock}:
c1, cm, c2 = ch[f], args[0], args[1]
args = [c1, cm, c2, *args[2:]]
if m is HGBlock:
args.insert(4, n) # number of repeats
n = 1
elif m is ResNetLayer:
c2 = args[1] if args[3] else args[1] * 4
elif m is nn.BatchNorm2d:
args = [ch[f]]
elif m is Concat:
c2 = sum(ch[x] for x in f)
elif m in {Detect, WorldDetect, Segment, Pose, OBB, ImagePoolingAttn}:
args.append([ch[x] for x in f])
if m is Segment:
args[2] = make_divisible(min(args[2], max_channels) * width, 8)
elif m is RTDETRDecoder: # special case, channels arg must be passed in index 1
args.insert(1, [ch[x] for x in f])
elif m is CBLinear:
c2 = args[0]
c1 = ch[f]
args = [c1, c2, *args[1:]]
elif m is CBFuse:
c2 = ch[f[-1]]
elif m is CBAM:
c1, c2 = ch[f], args[0]
if c2 != nc:
c2 = make_divisible(min(c2, max_channels)*width, 8)
args = [c1, *args[1:]]
else:
c2 = ch[f]
elif m is CBAM:
c1, c2 = ch[f], args[0]
if c2 != nc:
c2 = make_divisible(min(c2, max_channels)*width, 8)
args = [c1, *args[1:]]
查看conv.py模块是否存在“CBAM”
查看__init__.py文件中CBAM模块
修改配置文件,复制一份yolov8.yaml文件,命名位yolov8_CBAM.yaml文件
# YOLOv8.0n backbone
backbone:
# [from, repeats, module, args]
- [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
- [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
- [-1, 3, C2f, [128, True]]
# - [-1, 3, CBAM, [128, 7]] #
- [-1, 1, Conv, [256, 3, 2]] # 3-P3/8
- [-1, 6, C2f, [256, True]]
# - [-1, 6, CBAM, [256, 7]] #
- [-1, 1, Conv, [512, 3, 2]] # 5-P4/16
- [-1, 6, C2f, [512, True]]
# - [-1, 6, CBAM, [512, 7]] #
- [-1, 1, Conv, [1024, 3, 2]] # 7-P5/32
- [-1, 3, C2f, [1024, True]]
# - [-1, 3, CBAM, [1024, 7]] #
- [-1, 1, SPPF, [1024, 5]] #
# YOLOv8.0n head
head:
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 6], 1, Concat, [1]] # cat backbone P4
- [-1, 3, C2f, [512]] # 12
- [-1, 3, CBAM, [512, 7]]
- [-1, 1, nn.Upsample, [None, 2, "nearest"]]
- [[-1, 4], 1, Concat, [1]] # cat backbone P3
- [-1, 3, C2f, [256]] # (P3/8-small)
- [-1, 3, CBAM, [256, 7]]
- [-1, 1, Conv, [256, 3, 2]]
- [[-1, 12], 1, Concat, [1]] # cat head P4
- [-1, 3, C2f, [512]] # (P4/16-medium)
- [-1, 3, CBAM, [512, 7]]
- [-1, 1, Conv, [512, 3, 2]]
- [[-1, 9], 1, Concat, [1]] # cat head P5
- [-1, 3, C2f, [1024]] # (P5/32-large)
- [-1, 3, CBAM, [1024, 7]]
- [[16, 20, 24], 1, Detect, [nc]] # Detect(P3, P4, P5)
新建文件查看模块是否添加成功
from ultralytics import YOLO
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
molel = YOLO('ultralytics/cfg/models/v8/yolov8_CBAM_mutl.yaml', verbose=True)
添加成功