【YOLO改进】主干插入ParNetAttention模块(基于MMYOLO)

本文链接：https://blog.youkuaiyun.com/Vlone_pp/article/details/137523829

本文介绍了如何将ParNetAttention模块集成到MMYOLO框架中，包括修改PyTorch代码、注册模块、配置文件的调整，以及如何在训练过程中使用。

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

ParNetAttention模块

论文链接:https://arxiv.org/pdf/2110.07641.pdf

将ParNetAttention模块添加到MMYOLO中

将开源代码ParNetAttention.py文件复制到mmyolo/models/plugins目录下
导入MMYOLO用于注册模块的包: from mmyolo.registry import MODELS
确保 class ParNetAttention中的输入维度为in_channels(因为MMYOLO会提前传入输入维度参数，所以要保持参数名的一致)
利用@MODELS.register_module()将“class ParNetAttention(nn.Module)”注册:
修改mmyolo/models/plugins/__init__.py文件
在终端运行:
```
python setup.py install
```
修改对应的配置文件，并且将plugins的参数“type”设置为“ParNetAttention”，可参考【YOLO改进】主干插入注意力机制模块CBAM（基于MMYOLO）-优快云博客

修改后的ParNetAttention.py

import torch
from torch import nn
from mmyolo.registry import MODELS


@MODELS.register_module()
class ParNetAttention(nn.Module):

    def __init__(self, in_channels=512):
        super().__init__()
        self.sse = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, in_channels, kernel_size=1),
            nn.Sigmoid()
        )

        self.conv1x1 = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, kernel_size=1),
            nn.BatchNorm2d(in_channels)
        )
        self.conv3x3 = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(in_channels)
        )
        self.silu = nn.SiLU()

    def forward(self, x):
        b, c, _, _ = x.size()
        x1 = self.conv1x1(x)
        x2 = self.conv3x3(x)
        x3 = self.sse(x) * x
        y = self.silu(x1 + x2 + x3)
        return y


if __name__ == '__main__':
    input = torch.randn(50, 512, 7, 7)
    pna = ParNetAttention(channel=512)
    output = pna(input)
    print(output.shape)

修改后的init.py

# Copyright (c) OpenMMLab. All rights reserved.
from .cbam import CBAM
from .Biformer import BiLevelRoutingAttention
from .A2Attention import DoubleAttention
from .CoordAttention import CoordAtt
from .CoTAttention import CoTAttention
from .ECA import ECAAttention
from .EffectiveSE import EffectiveSEModule
from .EMA import EMA
from .GC import GlobalContext
from .GE import GatherExcite
from .MHSA import MHSA
from .ParNetAttention import ParNetAttention
__all__ = ['CBAM', 'BiLevelRoutingAttention', 'DoubleAttention', 'CoordAtt','CoTAttention','ECAAttention','EffectiveSEModule','EMA',
           'GlobalContext', 'GatherExcite', 'MHSA', 'ParNetAttention'