以下是在 YOLOv7 中添加 GAM 注意力机制的代码实现:

  1. 首先,需要在 model.py 中添加 GAM 注意力机制的类:
class GAM(nn.Module):
    def __init__(self, in_channels, out_channels, reduction=16):
        super(GAM, self).__init__()
        self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv1 = nn.Conv2d(in_channels, out_channels // reduction, kernel_size=1, bias=False)
        self.conv2 = nn.Conv2d(out_channels // reduction, out_channels, kernel_size=1, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.global_avg_pool(x)
        y = self.conv1(y)
        y = F.relu(y, inplace=True)
        y = self.conv2(y)
        y = self.sigmoid(y)
        return x * y.expand_as(x)
  1. 在 CSPDarknet 中添加 GAM 注意力机制:
class CSPBlock(nn.Module):
    def __init__(self, in_channels, out_channels, num_blocks, use_gam=True):
        super(CSPBlock, self).__init__()
        self.split = nn.Sequential(
            Conv(in_channels, out_channels, kernel_size=3, stride=2),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            *[Residual(out_channels, use_gam=use_gam) for _ in range(num_blocks)]
        )
        self.downsample = nn.Sequential(
            Conv(out_channels, out_channels, kernel_size=1, stride=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            Conv(out_channels, out_channels, kernel_size=3, stride=2),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            *[Residual(out_channels, use_gam=use_gam) for _ in range(num_blocks)],
        )
        self.merge = nn.Sequential(
            Conv(2 * out_channels, out_channels, kernel_size=1, stride=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            *[Residual(out_channels, use_gam=use_gam) for _ in range(num_blocks)]
        )

    def forward(self, x):
        x1 = self.split(x)
        x2 = self.downsample(x1)
        x = torch.cat([x1, x2], dim=1)
        x = self.merge(x)
        return x

在 Residual 模块中添加 GAM 注意力机制:

class Residual(nn.Module):
    def __init__(self, channels, use_gam=True):
        super(Residual, self).__init__()
        self.use_gam = use_gam
        self.conv1 = Conv(channels, channels, kernel_size=1, stride=1)
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = Conv(channels, channels, kernel_size=3, stride=1)
        self.bn2 = nn.BatchNorm2d(channels)
        self.conv3 = Conv(channels, channels, kernel_size=1, stride=1)
        self.bn3 = nn.BatchNorm2d(channels)
        if self.use_gam:
            self.gam = GAM(channels, channels)

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = F.relu(out, inplace=True)
        out = self.conv2(out)
        out = self.bn2(out)
        out = F.relu(out, inplace=True)
        out = self.conv3(out)
        out = self.bn3(out)
        if self.use_gam:
            out = self.gam(out)
        out += residual
        out = F.relu(out, inplace=True)
        return out
  1. 在 YOLOv7 中使用 CSPDarknet with GAM 注意力机制:
class YOLOv7(nn.Module):
    def __init__(self, num_classes=80):
        super(YOLOv7, self).__init__()
        self.backbone = CSPDarknet(use_gam=True)
        self.neck = SPP(use_gam=True)
        self.head = YOLOHead(in_channels=[128, 256, 512], num_anchors=3, num_classes=num_classes)

    def forward(self, x):
        x = self.backbone(x)
        x = self.neck(x)
        x = self.head(x)
        return x

这样,就在 YOLOv7 中添加了 GAM 注意力机制。

如何在 YOLOv7 中添加 GAM 注意力机制:代码实现

原文地址: https://www.cveoy.top/t/topic/m4mC 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录