Lightweight Hybrid Pooling Attention Mechanism: Plug-and-Play for Enhanced Performance

Lightweight Hybrid Pooling Attention Mechanism: Plug-and-Play for Enhanced Performance

Title:A synergistic CNN-transformer network with pooling attention fusion for hyperspectral image classification

Paper link: https://github.com/chenpeng052/SCT-Net

Lightweight Hybrid Pooling Attention Mechanism: Plug-and-Play for Enhanced Performance

Lightweight Hybrid Pooling Attention Mechanism: Plug-and-Play for Enhanced Performance

  • Proposed a CNN-Transformer synergistic network, combining the local feature extraction capability of CNNs with the global modeling capability of Transformers.

  • Designed a Twin-Branch Feature Extraction (TBFE) module, which uses 2D and 3D convolutions in parallel to extract spatial and spectral features simultaneously.

  • Introduced a Hybrid Pooling Attention (HPA) module, which combines average pooling and max pooling to achieve global dependency modeling across space.

  • Designed a Cross-layer Feature Fusion (CFF) module to alleviate the information loss problem between different Transformer encoding layers.

  • Validated on multiple public hyperspectral datasets, significantly outperforming existing methods.

Lightweight Hybrid Pooling Attention Mechanism: Plug-and-Play for Enhanced Performance

Overall Structure

The model first reduces the dimensionality of hyperspectral images using PCA and segments them into patches, then utilizes the TBFE Twin-Branch Convolution Module to extract spectral and spatial features simultaneously, followed by the HPA module to enhance channel and spatial attention, and then inputs into the Transformer encoder for global modeling. Finally, the CFF Cross-layer Fusion alleviates inter-layer information loss and completes pixel-level classification through a fully connected layer, achieving an organic combination of CNN’s local perception and Transformer’s global dependency modeling.Lightweight Hybrid Pooling Attention Mechanism: Plug-and-Play for Enhanced Performance

  • Data Preprocessing: PCA is used for dimensionality reduction on hyperspectral data, extracting 3D patches.

  • Twin-Branch Feature Extraction (TBFE): Uses 2D convolution for the spatial branch and 3D convolution for the spectral branch to extract features in parallel, merging to obtain a joint feature representation.

  • Hybrid Pooling Attention (HPA): Utilizes average pooling and max pooling to establish attention in horizontal and vertical directions, weighted fusion of cross-channel and cross-space information.

  • Vision Transformer Encoder: Performs global modeling on features extracted by HPA to obtain high-level spectral-spatial features.

  • Cross-layer Feature Fusion (CFF): Conducts cross-layer fusion of Transformer encoding results to reduce information loss.

  • Classification Head: Outputs the class label of the central pixel through a fully connected layer.

Lightweight Hybrid Pooling Attention Mechanism: Plug-and-Play for Enhanced Performance

HPA is a lightweight cross-space-channel attention module suitable for insertion into CNNs or Transformers, designed to enhance feature representation in scenarios such as remote sensing and medical imaging, highlighting key areas and improving classification or recognition accuracy.

1. Applicable Scenarios

  • Hyperspectral Image Classification / Remote Sensing Scenarios: Requires capturing both local and global dependencies across channels and space.

  • Deep Models Requiring Enhanced Feature Representation: For example, convolutional networks or Transformers, HPA can be inserted as a lightweight attention module.

  • Few-Shot Learning: When samples are limited, HPA can enhance feature utilization through pooling, reducing overfitting.

  • Multimodal Feature Fusion: Suitable for tasks requiring interaction between different channels or spatial scales (e.g., medical imaging, video understanding).

2. Functions Performed

  • Cross-Channel Enhancement: By modeling with parallel average pooling and max pooling, it retains overall trends while capturing significant feature points, achieving adaptive weighting of channel attention.

  • Cross-Space Dependency Modeling: Performs 1D global pooling in both horizontal and vertical directions, supplementing the shortcomings of Transformers or CNNs in long-distance dependencies.

  • Improving Discriminability: By re-weighting feature maps, it highlights key areas (e.g., object boundaries or small targets) while suppressing redundant information.

  • Lightweight and Efficient: Compared to standard self-attention, HPA has fewer parameters and lower computational overhead, making it easy to embed as a module in other networks.

Lightweight Hybrid Pooling Attention Mechanism: Plug-and-Play for Enhanced Performance

import torch
from torch import nn

# Hybrid pooling attention (HPA)
class HPA(nn.Module):
    def __init__(self, channels, c2=None, factor=32):
        super(HPA, self).__init__()
        self.groups = factor
        assert channels // self.groups > 0
        self.softmax = nn.Softmax(-1)
        self.agp = nn.AdaptiveAvgPool2d((1, 1))
        self.map = nn.AdaptiveMaxPool2d((1, 1))
        self.pool_h = nn.AdaptiveAvgPool2d((None, 1)) #Y avg
        self.pool_w = nn.AdaptiveAvgPool2d((1, None)) #X avg
        self.max_h = nn.AdaptiveMaxPool2d((None, 1)) #Y avg
        self.max_w = nn.AdaptiveMaxPool2d((1, None)) #X avg

        self.gn = nn.GroupNorm(channels // self.groups, channels // self.groups)
        self.conv1x1 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=1, stride=1, padding=0)
        self.conv3x3 = nn.Conv2d(channels // self.groups, channels // self.groups, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        b, c, h, w = x.size()
        group_x = x.reshape(b * self.groups, -1, h, w) # b*g,c//g,h,w --->2048,2,11,11
        x_h = self.pool_h(group_x) #2048,2,11,1
        x_w = self.pool_w(group_x).permute(0, 1, 3, 2) #2048,2,1,11--->2048,2,11,1
        hw = self.conv1x1(torch.cat([x_h, x_w], dim=2)) #2048,2,22,1
        x_h, x_w = torch.split(hw, [h, w], dim=2) #2048,2,11,1
        x1 = self.gn(group_x * x_h.sigmoid() * x_w.permute(0, 1, 3, 2).sigmoid()) #2048,2,11,11
        x2 = self.conv3x3(group_x) #2048,2,11,11

        y_h = self.max_h(group_x) #2048,2,11,1
        y_w = self.max_w(group_x).permute(0, 1, 3, 2)
        yhw = self.conv1x1(torch.cat([y_h, y_w], dim=2)) #2048,2,22,1
        y_h, y_w = torch.split(yhw, [h, w], dim=2) #2048,2,11,1
        y1 = self.gn(group_x * y_h.sigmoid() * y_w.permute(0, 1, 3, 2).sigmoid()) #2048,2,11,11
        y11 = y1.reshape(b * self.groups, c // self.groups, -1) # b*g, c//g, hw 2048,2,121
        y12 = self.softmax(self.map(y1).reshape(b * self.groups, -1, 1).permute(0, 2, 1)) #2048,1,2

        x11 = x1.reshape(b * self.groups, c // self.groups, -1) # b*g, c//g, hw 2048,2,121
        x12 = self.softmax(self.agp(x1).reshape(b * self.groups, -1, 1).permute(0, 2, 1)) #2048,2,1,1-->2048,2,1--->2048,1,2
        x21 = x2.reshape(b * self.groups, c // self.groups, -1) # b*g, c//g, hw #2048,2,121
        x22 = self.softmax(self.agp(x2).reshape(b * self.groups, -1, 1).permute(0, 2, 1)) #2048,2,1,1-->2048,2,1--->2048,1,2
        weights = (torch.matmul(x12, y11) + torch.matmul(y12, x11)).reshape(b * self.groups, 1, h, w)
        return (group_x * weights.sigmoid()).reshape(b, c, h, w) 
    
if __name__ == "__main__":
    # Module parameters
    batch_size = 1    # Batch size
    channels = 32     # Number of input feature channels
    height = 256      # Image height
    width = 256        # Image width

    model = HPA(channels=channels)
    print(model)

    # Generate random input tensor (batch_size, channels, height, width)
    x = torch.randn(batch_size, channels, height, width)

    # Print the shape of the input tensor
    print("Input shape:", x.shape)

    # Forward propagation to compute output
    output = model(x)

    # Print the shape of the output tensor
    print("Output shape:", output.shape)

For more analysis, see the original text.

Lightweight Hybrid Pooling Attention Mechanism: Plug-and-Play for Enhanced Performance

Leave a Comment