【YOLO改进】主干插入注意力机制模块CBAM(基于MMYOLO)

本文介绍了如何在MMYOLO模型中插入CBAM注意力机制模块,通过修改配置文件,如yolov5_s-v61_syncbn_8xb16-300e_coco.py,以增强模型的特征提取能力,并可自定义插件的插入位置和类型。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

MMYOLO插入CBAM 

由于CBAM是MMYOLO自带的注意力机制模块,所以在插入CBAM的时候直接修改对应的配置文件即可,以 configs/yolov5/yolov5_s-v61_syncbn_8xb16-300e_coco.py为例子,将model里面的backbone修改即可:

model = dict(
    backbone=dict(frozen_stages=4,
    ## 在原有设置的基础上,加上插件的设置
    plugins=[
              dict(cfg=dict(type='CBAM'),
                     stages=(False, True, True, True))
            ],
                  ),
    bbox_head=dict(
        head_module=dict(num_classes=num_classes),
        prior_generator=dict(base_sizes=anchors)))

在插件的设置中, 可以通过修改stages的参数自定义模块插入在主干的哪一个stage,也可以通过修改“type”参数选择其他注意力机制模块。一个简单的例子如下:

        Examples:
            >>> plugins=[
            ...     dict(cfg=dict(type='xxx', arg1='xxx'),
            ...          stages=(False, True, True, True)),
            ...     dict(cfg=dict(type='yyy'),
            ...          stages=(True, True, True, True)),
            ... ]
        Suppose ``stage_idx=0``, the structure of blocks in the stage would be:

        .. code-block:: none

            conv1 -> conv2 -> conv3 -> yyy

        Suppose ``stage_idx=1``, the structure of blocks in the stage would be:

        .. code-block:: none

            conv1 -> conv2 -> conv3 -> xxx -> yyy

由于CBAM是MMYOLO自带的,所以直接设置参数“type”为“CBAM”即可,如果需要用其他注意力机制模块,则需要通过注册自定义模块实现。

修改后的configs/yolov5/yolov5_s-v61_syncbn_8xb16-300e_coco.py

_base_ = ['../_base_/default_runtime.py', '../_base_/det_p5_tta.py']

# ========================Frequently modified parameters======================
# -----data related-----
data_root = 'data/coco/'  # Root path of data
# Path of train annotation file
train_ann_file = 'annotations/instances_train2017.json'
train_data_prefix = 'train2017/'  # Prefix of train image path
# Path of val annotation file
val_ann_file = 'annotations/instances_val2017.json'
val_data_prefix = 'val2017/'  # Prefix of val image path

num_classes = 80  # Number of classes for classification
# Batch size of a single GPU during training
train_batch_size_per_gpu = 16
# Worker to pre-fetch data for each single GPU during training
train_num_workers = 8
#
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值