stevengrove
initial commit
186701e
|
raw
history blame
1.43 kB

给主干网络增加插件

MMYOLO 支持在 Backbone 的不同 Stage 后增加如 none_localdropblock 等插件,用户可以直接通过修改 config 文件中 backboneplugins参数来实现对插件的管理。例如为 YOLOv5 增加GeneralizedAttention 插件,其配置文件如下:

_base_ = './yolov5_s-v61_syncbn_8xb16-300e_coco.py'

model = dict(
    backbone=dict(
        plugins=[
            dict(
                cfg=dict(
                    type='GeneralizedAttention',
                    spatial_range=-1,
                    num_heads=8,
                    attention_type='0011',
                    kv_stride=2),
                stages=(False, False, True, True))
        ]))

cfg 参数表示插件的具体配置, stages 参数表示是否在 backbone 对应的 stage 后面增加插件,长度需要和 backbone 的 stage 数量相同。

目前 MMYOLO 支持了如下插件:

支持的插件
  1. CBAM
  2. GeneralizedAttention
  3. NonLocal2d
  4. ContextBlock