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

TripletAttention模块

论文链接:https://arxiv.org/abs/2010.03045

将TripletAttention模块添加到MMYOLO中

1. 将开源代码TripletAttention.py文件复制到mmyolo/models/plugins目录下

2. 导入MMYOLO用于注册模块的包: from mmyolo.registry import MODELS

3. 确保 class TripletAttention中的输入维度为in_channels(因为MMYOLO会提前传入输入维度参数，所以要保持参数名的一致)

4. 利用@MODELS.register_module()将“class TripletAttention(nn.Module)”注册:

5. 修改mmyolo/models/plugins/__init__.py文件

6. 在终端运行:

   python setup.py install

7. 修改对应的配置文件，并且将plugins的参数“type”设置为“TripletAttention”，可参考【YOLO改进】主干插入注意力机制模块CBAM（基于MMYOLO）-优快云博客

修改后的TripletAttention.py

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



class BasicConv(nn.Module):
    def __init__(self, in_planes, out_planes, kernel_size, stride=1, padding=0, dilation=1, groups=1, relu=True,
                 bn=True, bias=False):
        super(BasicConv, self).__init__()
        self.out_channels = out_planes
        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=padding,
                              dilation=dilation, groups=groups, bias=bias)
        self.bn = nn.BatchNorm2d(out_planes, eps=1e-5, momentum=0.01, affine=True) if bn else None
        self.relu = nn.ReLU() if relu else None

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.relu is not None:
            x = self.relu(x)
        return x


class ZPool(nn.Module):
    def forward(self, x):
        return torch.cat((torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1)


class AttentionGate(nn.Module):
    def __init__(self):
        super(AttentionGate, self).__init__()
        kernel_size = 7
        self.compress = ZPool()
        self.conv = BasicConv(2, 1, kernel_size, stride=1, padding=(kernel_size - 1) // 2, relu=False)

    def forward(self, x):
        x_compress = self.compress(x)
        x_out = self.conv(x_compress)
        scale = torch.sigmoid_(x_out)
        return x * scale

@MODELS.register_module()
class TripletAttention(nn.Module):
    def __init__(self, in_channels, no_spatial=False):
        super(TripletAttention, self).__init__()
        self.cw = AttentionGate()
        self.hc = AttentionGate()
        self.no_spatial = no_spatial
        if not no_spatial:
            self.hw = AttentionGate()

    def forward(self, x):
        x_perm1 = x.permute(0, 2, 1, 3).contiguous()
        x_out1 = self.cw(x_perm1)
        x_out11 = x_out1.permute(0, 2, 1, 3).c