[YOLOv5-6.x]CAアテンションメカニズムを追加してもFLOPが表示されないという問題を解決します

1.問題の説明

問題は以前に書かれたブログから来ています：[魔法の変更YOLOv5-6.x（中央）]：ACON活性化関数、CBAMおよびCAアテンションメカニズム、加重双方向機能ピラミッドBiFPNを追加し、YOLOv5のバックボーンに座標アテンションを追加してみてください。参加後、mAPは改善されましたが、重要なGFLOP情報は表示できません。！

2.問題解決

それで、私はクレイジーなデバッグを開始し、さまざまな場所にブレークポイントを追加し、ブレークポイントを追加し、最後にソースをロックしました：CABlockクラス関数の定義

これはCABlockの元のクラス関数であり、__init__2つの適応平均プーリングが定義されnn.AdaptiveAvgPool2d()ています。CAを追加した後にGFLOP情報が表示されない原因は、これら2つの適応平均プーリングです。具体的な理由はまだ不明です。皆さん、ようこそ。コミュニケーション〜

class CABlock(nn.Module):
    def __init__(self, inp, oup, reduction=32):
        super(CABlock, self).__init__()
        # ============================
        # height方向上的均值池化
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1))
        # width方向上的均值池化
        self.pool_w = nn.AdaptiveAvgPool2d((1, None))

        mip = max(8, inp // reduction)

        self.conv1 = nn.Conv2d(inp, mip, kernel_size=1, stride=1, padding=0)
        self.bn1 = nn.BatchNorm2d(mip)
        self.act = h_swish()

        self.conv_h = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)
        self.conv_w = nn.Conv2d(mip, oup, kernel_size=1, stride=1, padding=0)

    def forward(self, x):
        identity = x

        n, c, h, w = x.size()
        x_h = self.pool_h(x)
        x_w = self.pool_w(x).permute(0, 1, 3, 2)

        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.act(y)

        x_h, x_w = torch.split(y, [h, w], dim=2)
        x_w = x_w.permute(0, 1, 3, 2)

        a_h = self.conv_h(x_h).sigmoid()
        a_w = self.conv_w(x_w).sigmoid()

        out = identity * a_w * a_h

        return out

そこで、以下に変更して問題を解決しました

class CABlock(nn.Module):
    def __init__(self, inp, oup, reduction=32):
        super(CABlock, self).__init__()
        mip = max(8, inp // reduction)

        self.conv1 = nn.Conv2d(inp, mip, 1, 1, bias=False)
        self.bn1 = nn.BatchNorm2d(mip)
        self.act = h_swish()

        self.conv_h = nn.Conv2d(mip, oup, 1, 1, bias=False)
        self.conv_w = nn.Conv2d(mip, oup, 1, 1, bias=False)

    def forward(self, x):
        identity = x
        _, _, h, w = x.size()
        pool_h = nn.AdaptiveAvgPool2d((h, 1))
        x_h = pool_h(x)
        pool_w = nn.AdaptiveAvgPool2d((1, w))
        x_w = pool_w(x).permute(0, 1, 3, 2)

        y = torch.cat([x_h, x_w], dim=2)
        y = self.conv1(y)
        y = self.bn1(y)
        y = self.act(y)

        y_h, y_w = torch.split(y, [h, w], dim=2)
        y_w = y_w.permute(0, 1, 3, 2)

        a_h = self.conv_h(y_h).sigmoid()
        a_w = self.conv_w(y_w).sigmoid()

        return identity * a_w * a_h