YOLO即插即用---PKIBlock

Poly Kernel Inception Network for Remote Sensing Detection

论文地址

1. 解决的问题

2. 解决方案

3. 解决问题的具体方法

4. 模块的应用

5. 在目标检测任务中的添加位置

6.即插即用代码

论文地址

2403.06258https://arxiv.org/pdf/2403.06258

1. 解决的问题

遥感图像目标检测面临着两大挑战：

目标尺度变化大: 遥感图像通常包含各种尺度差异巨大的目标，例如足球场、建筑物、车辆等。这种尺度变化对目标检测算法提出了挑战，因为传统的目标检测方法往往难以有效地处理不同尺度的目标。
场景背景复杂多样: 遥感图像通常包含复杂的背景信息，例如山脉、河流、植被等。这些背景信息会对目标检测造成干扰，使得算法难以准确地识别和定位目标。

2. 解决方案

PKINet 通过以下两种机制来解决上述挑战：

多尺度卷积核 (PKI 模块):
- PKI 模块使用并行排列的不同尺寸深度可分离卷积核，提取不同感受野的多尺度纹理特征。这种设计可以有效地捕捉到不同尺度目标的特征，从而提高模型对不同尺度目标的检测能力。
- 深度可分离卷积可以有效地降低模型的计算量和参数量，使得模型更加轻量化。
上下文锚点注意力机制 (CAA 模块):
- CAA 模块利用全局平均池化和 1D 线条卷积来捕捉远距离像素之间的关系，并增强中心区域的特征。这种设计可以帮助模型更好地理解目标的上下文信息，从而提高模型的检测精度。
- 1D 线条卷积可以有效地捕捉长距离像素关系，而不会引入过多的噪声和计算量。

3. 解决问题的具体方法

PKI 模块:
- 局部特征提取: 使用 3x3 卷积提取局部特征，捕获目标的细节信息。
- 多尺度上下文特征提取: 使用不同尺寸 (3x3, 5x5, 7x7, 9x9, 11x11) 的深度可分离卷积提取不同感受野的上下文特征，捕捉目标周围的环境信息。
- 特征融合: 使用 1x1 卷积融合局部特征和上下文特征，得到最终的特征表示。
CAA 模块:
- 局部区域特征提取: 使用全局平均池化和 1x1 卷积提取局部区域特征，作为后续计算的基础。
- 远距离像素关系捕捉: 使用 1D 线条卷积 (水平方向和垂直方向) 代替大核卷积，扩大感受野并捕捉远距离像素关系，从而更好地理解目标的上下文信息。
- 注意力权重生成: 使用 sigmoid 函数生成注意力权重，并根据注意力权重增强 PKI 模块的输出特征，突出目标区域的特征，抑制背景区域的特征。

4. 模块的应用

PKI 模块和 CAA 模块共同构成了 PKINet 的基本单元，用于提取图像特征。PKINet 可以作为特征提取骨干网络，应用于各种遥感图像目标检测任务，例如检测飞机、舰船、车辆等目标。通过 PKINet 提取的特征可以更好地反映目标的尺度、形状、纹理和上下文信息，从而提高目标检测算法的性能。

5. 在目标检测任务中的添加位置

PKINet 位于目标检测网络的底层，负责提取图像特征。这些特征随后被送入检测头，例如 RPN、RoI Pooling、分类和回归头等，用于识别和定位目标。PKINet 提取的特征可以帮助检测头更好地理解目标的特征，从而提高目标检测算法的精度和鲁棒性

6.即插即用代码

from typing import Optional, Sequence
import torch.nn as nn
import torch
def autopad(kernel_size: int, padding: Optional[int] = None, dilation: int = 1) -> int:
    """Calculate the padding size based on kernel size and dilation."""
    if padding is None:
        padding = (kernel_size - 1) * dilation // 2
    return padding


def make_divisible(value: int, divisor: int = 8) -> int:
    """Make a value divisible by a certain divisor."""
    return int((value + divisor // 2) // divisor * divisor)


class ConvModule(nn.Module):
    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            kernel_size: int,
            stride: int = 1,
            padding: int = 0,
            dilation: int = 1,
            groups: int = 1,
            norm_cfg: Optional[dict] = None,
            act_cfg: Optional[dict] = None):
        super().__init__()
        layers = []
        layers.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation=dilation, groups=groups, bias=(norm_cfg is None)))
        if norm_cfg:
            norm_layer = self._get_norm_layer(out_channels, norm_cfg)
            layers.append(norm_layer)
        if act_cfg:
            act_layer = self._get_act_layer(act_cfg)
            layers.append(act_layer)
        self.block = nn.Sequential(*layers)

    def forward(self, x):
        return self.block(x)

    def _get_norm_layer(self, num_features, norm_cfg):
        if norm_cfg['type'] == 'BN':
            return nn.BatchNorm2d(num_features, momentum=norm_cfg.get('momentum', 0.1), eps=norm_cfg.get('eps', 1e-5))
        # Add more normalization types if needed
        raise NotImplementedError(f"Normalization layer '{norm_cfg['type']}' is not implemented.")

    def _get_act_layer(self, act_cfg):
        if act_cfg['type'] == 'ReLU':
            return nn.ReLU(inplace=True)
        if act_cfg['type'] == 'SiLU':
            return nn.SiLU(inplace=True)
        # Add more activation types if needed
        raise NotImplementedError(f"Activation layer '{act_cfg['type']}' is not implemented.")

# Update InceptionBottleneck's constructor call to avoid conflicts
class InceptionBottleneck(nn.Module):
    """Bottleneck with Inception module"""
    def __init__(
            self,
            in_channels: int,
            out_channels: Optional[int] = None,
            kernel_sizes: Sequence[int] = (3, 5, 7, 9, 11),
            dilations: Sequence[int] = (1, 1, 1, 1, 1),
            expansion: float = 1.0,
            add_identity: bool = True,
            with_caa: bool = True,
            caa_kernel_size: int = 11,
            norm_cfg: Optional[dict] = dict(type='BN', momentum=0.03, eps=0.001),
            act_cfg: Optional[dict] = dict(type='SiLU')):
        super().__init__()
        out_channels = out_channels or in_channels
        hidden_channels = make_divisible(int(out_channels * expansion), 8)

        self.pre_conv = ConvModule(in_channels, hidden_channels, 1, 1, 0,
                                   norm_cfg=norm_cfg, act_cfg=act_cfg)

        self.dw_conv = ConvModule(hidden_channels, hidden_channels, kernel_sizes[0], 1,
                                  autopad(kernel_sizes[0], None, dilations[0]),
                                  dilation=dilations[0], groups=hidden_channels,
                                  norm_cfg=None, act_cfg=None)
        self.dw_conv1 = ConvModule(hidden_channels, hidden_channels, kernel_sizes[1], 1,
                                   autopad(kernel_sizes[1], None, dilations[1]),
                                   dilation=dilations[1], groups=hidden_channels,
                                   norm_cfg=None, act_cfg=None)
        self.dw_conv2 = ConvModule(hidden_channels, hidden_channels, kernel_sizes[2], 1,
                                   autopad(kernel_sizes[2], None, dilations[2]),
                                   dilation=dilations[2], groups=hidden_channels,
                                   norm_cfg=None, act_cfg=None)
        self.dw_conv3 = ConvModule(hidden_channels, hidden_channels, kernel_sizes[3], 1,
                                   autopad(kernel_sizes[3], None, dilations[3]),
                                   dilation=dilations[3], groups=hidden_channels,
                                   norm_cfg=None, act_cfg=None)
        self.dw_conv4 = ConvModule(hidden_channels, hidden_channels, kernel_sizes[4], 1,
                                   autopad(kernel_sizes[4], None, dilations[4]),
                                   dilation=dilations[4], groups=hidden_channels,
                                   norm_cfg=None, act_cfg=None)
        self.pw_conv = ConvModule(hidden_channels, hidden_channels, 1, 1, 0,
                                  norm_cfg=norm_cfg, act_cfg=act_cfg)

        if with_caa:
            self.caa_factor = CAA(hidden_channels, caa_kernel_size, caa_kernel_size, None, None)
        else:
            self.caa_factor = None

        self.add_identity = add_identity and in_channels == out_channels

        self.post_conv = ConvModule(hidden_channels, out_channels, 1, 1, 0,
                                    norm_cfg=norm_cfg, act_cfg=act_cfg)

    def forward(self, x):
        x = self.pre_conv(x)

        y = x
        x = self.dw_conv(x)
        x = x + self.dw_conv1(x) + self.dw_conv2(x) + self.dw_conv3(x) + self.dw_conv4(x)
        x = self.pw_conv(x)
        if self.caa_factor is not None:
            y = self.caa_factor(y)
        if self.add_identity:
            y = x * y
            x = x + y
        else:
            x = x * y

        x = self.post_conv(x)
        return x

class CAA(nn.Module):
    """Context Anchor Attention"""
    def __init__(
            self,
            channels: int,
            h_kernel_size: int = 11,
            v_kernel_size: int = 11,
            norm_cfg: Optional[dict] = dict(type='BN', momentum=0.03, eps=0.001),
            act_cfg: Optional[dict] = dict(type='SiLU')):
        super().__init__()
        self.avg_pool = nn.AvgPool2d(7, 1, 3)
        self.conv1 = ConvModule(channels, channels, 1, 1, 0,
                                norm_cfg=norm_cfg, act_cfg=act_cfg)
        self.h_conv = ConvModule(channels, channels, (1, h_kernel_size), 1,
                                 (0, h_kernel_size // 2), groups=channels,
                                 norm_cfg=None, act_cfg=None)
        self.v_conv = ConvModule(channels, channels, (v_kernel_size, 1), 1,
                                 (v_kernel_size // 2, 0), groups=channels,
                                 norm_cfg=None, act_cfg=None)
        self.conv2 = ConvModule(channels, channels, 1, 1, 0,
                                norm_cfg=norm_cfg, act_cfg=act_cfg)
        self.act = nn.Sigmoid()

    def forward(self, x):
        attn_factor = self.act(self.conv2(self.v_conv(self.h_conv(self.conv1(self.avg_pool(x))))))
        return attn_factor

# Testing the InceptionBottleneck
if __name__ == "__main__":

    input = torch.randn(1, 64, 128, 128) #输入B C H W
    block = InceptionBottleneck(in_channels=64, out_channels=128)
    output = block(input)
    print(input.size())
    print(output.size())

对模型改进感兴趣的可以进入交流群，群中有答疑（QQ：828370883）