Mecanismo de atenção - Módulo de atenção de bloqueio convolucional (CBAM)

Módulo de Atenção de Bloco Convolucional (CBAM) : CBAM é um modelo de combinação que combina atenção de canal e atenção espacial para melhorar a expressividade do modelo.

O módulo CBAM inclui dois submódulos de atenção: módulo de atenção de canal e módulo de atenção espacial. O módulo de atenção do canal é usado para calcular a importância de cada canal para melhor distinguir os recursos entre os diferentes canais. O módulo de atenção espacial é usado para calcular a importância espacial de cada pixel, a fim de capturar melhor a estrutura espacial na imagem.

O módulo de atenção espacial executa agrupamento médio e agrupamento máximo na dimensão do canal do mapa de recursos de entrada e, em seguida, insere os dois resultados de agrupamento em uma camada totalmente conectada e, finalmente, gera um tensor de peso de atenção espacial. Esse tensor é usado para ponderar espacialmente cada pixel para capturar melhor a estrutura espacial na imagem.

A estrutura geral do módulo CBAM é mostrada na figura abaixo:

Na figura, as caixas verdes representam os módulos de atenção do canal e as caixas laranja representam os módulos de atenção espacial. Ao concatenar esses dois módulos, um módulo CBAM completo pode ser obtido para inserção em uma rede neural convolucional para melhorar o desempenho do modelo.

Implemente o CBAM com o pytorch:

import torch
import torch.nn as nn
import torch.nn.functional as F

class ChannelAttention(nn.Module):
    def __init__(self, in_channels, reduction_ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc1 = nn.Conv2d(in_channels, in_channels // reduction_ratio, 1, bias=False)
        self.relu = nn.ReLU()
        self.fc2 = nn.Conv2d(in_channels // reduction_ratio, in_channels, 1, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        x = self.conv(x)
        return self.sigmoid(x)


class CBAM(nn.Module):
    def __init__(self, in_channels, reduction_ratio=16, kernel_size=7):
        super(CBAM, self).__init__()
        self.channel_att = ChannelAttention(in_channels, reduction_ratio)
        self.spatial_att = SpatialAttention(kernel_size)

    def forward(self, x):
        out = self.channel_att(x) * x
        out = self.spatial_att(out) * out
        return out

O código acima implementa dois submódulos do módulo CBAM: o módulo de atenção do canal (ChannelAttention) e o módulo de atenção espacial (SpatialAttention), e todo o módulo CBAM (CBAM).

O módulo de atenção do canal executa o agrupamento máximo e o agrupamento médio na dimensão do canal do mapa de recursos de entrada e, em seguida, insere os dois resultados do agrupamento em uma camada totalmente conectada e, finalmente, gera um vetor de peso de atenção do canal. O módulo de atenção espacial executa agrupamento médio e agrupamento máximo na dimensão do canal do mapa de recursos de entrada e, em seguida, insere os dois resultados de agrupamento em uma camada totalmente conectada e, finalmente, gera um tensor de peso de atenção espacial. O módulo CBAM conecta esses dois submódulos em série para inserção na rede neural convolucional para melhorar o desempenho do modelo.

Use o módulo CBAM em seu modelo:

class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.cbam1 = CBAM(64)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.cbam2 = CBAM(128)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.cbam3 = CBAM(256)
        
        self.fc = nn.Linear(256 * 8 * 8, 10)
        
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.cbam1(x)
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        x = F.relu(self.conv2(x))
        x = self.cbam2(x)
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        x = F.relu(self.conv3(x))
        x = self.cbam3(x)
        x = F.max_pool2d(x, kernel_size=2, stride=2)
        
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        
        return x

Esta é uma rede neural convolucional simples com três camadas convolucionais e três módulos CBAM e, finalmente, uma camada totalmente conectada para converter o mapa de recursos em um rótulo previsto. No processo de passagem direta, o mapa de recursos de entrada é extraído e aprimorado por meio da camada convolucional e do módulo CBAM, em seguida, o tamanho espacial do mapa de recursos é reduzido por meio da camada de agrupamento máximo e, finalmente, o rótulo previsto é emitido por meio do totalmente conectado camada.

Mecanismo de atenção - Módulo de atenção de bloqueio convolucional (CBAM)

Acho que você gosta