OverLoCK：先概览，再聚焦。CVPR2025全新主干网络-EW帮帮网

1.简介

OverLoCK是一种创新的卷积神经网络(ConvNet)架构，它引入了概览优先（Overview-first）和聚焦次之（Look-Closely-next）的设计理念，通过上下文混合动态核(Context-Mixing Dynamic Kernels)实现了对图像特征的高效提取。该网络特别强调了在不同尺度上利用大小内核相结合的方式，以更好地捕捉多尺度特征，并通过一个称为ContMix的核心组件来实现长距离依赖关系和局部细节的建模。

实验结果显示，无论是在图像分类还是目标检测和实例分割任务中，OverLoCK都表现出了超越其他方法的性能。例如，在ImageNet-1K数据集上的图像分类任务中，OverLoCK模型达到了非常高的准确率。而在COCO 2017数据集上的目标检测和实例分割任务中，使用Mask R-CNN或Cascade Mask R-CNN框架时，OverLoCK-S相较于其他竞争对手如BiFormer-B和MogaNet-B，在APb指标上分别提升了0.8%和1.5%。此外，尽管基于ConvNet的方法在图像分类任务上与Transformer方法的表现相当，但在检测任务中却展示了显著的性能优势。

总之，OverLoCK不仅提供了对于传统卷积操作的有效改进，还为解决复杂的视觉识别任务提供了一种新的视角，特别是在增强特征提取能力和提高模型效率方面显示了巨大的潜力。

论文地址：[2502.20087] OverLoCK: An Overview-first-Look-Closely-next ConvNet with Context-Mixing Dynamic Kernels

github地址：https://github.com/LMMMEng/OverLoCK

2.论文详解

概述

作者提出了一种深度阶段分解策略（Deep-stage Decomposition Strategy, DDS），该策略受人类视觉系统“先概览后细看”（Overview-first-Look-Closely-next）机制的启发，将卷积网络（ConvNet）分解为三个协同工作的子网络：Base-Net（基础网络）、Overview-Net（概览网络）和Focus-Net（聚焦网络）。这种设计旨在通过动态的自上而下的语义引导来增强特征图和卷积核权重，从而提高模型性能。

网络结构

Base-Net：负责编码低层次和中层次的信息，通过三个嵌入层将输入图像逐步下采样至 H×W 的特征图。该特征图同时输入到Overview-Net和Focus-Net。
Overview-Net：是一个轻量级网络，快速生成一个语义上有意义但质量较低的全局上下文表示（即“概览”）。它通过立即下采样将中层次特征图压缩至 H/2×W/2，并作为上下文先验（context prior）反馈给Focus-Net。
Focus-Net：是一个更强大、更深的网络，它利用上下文先验来指导对中层次特征图的逐步细化，并扩大感受野，以获得更准确和信息丰富的高层次表示（即“细看”）。

Base-Net的输出特征图被分为两部分，一部分输入到Overview-Net，另一部分输入到Focus-Net。Focus-Net在每个块中将来自Overview-Net的上下文先验与自身的特征图融合，以实现特征级和权重级的引导。

其中的核心部分是GDSA中的ContMix。

ContMix

ContMix是一种新型的动态卷积方法——Context-Mixing Dynamic Convolution（ContMix）。该方法旨在使卷积操作能够同时建模长距离依赖关系，同时保持强大的局部归纳偏差。这是通过将全局上下文信息注入到每个卷积核权重中实现的，从而在固定大小的卷积核下也能有效地捕捉长距离依赖。

ContMix的核心思想是通过计算输入特征图中每个token与一组区域中心之间的亲和力值来表示token与上下文之间的关系。这些亲和力值被聚合起来以定义每个token的动态卷积核，从而将上下文知识注入到每个卷积核权重中。在卷积操作中，每个token都会与通过区域中心聚集的全局信息进行交互，从而实现长距离依赖的建模。

模块流程：

Token-wise Global Context Representation（全局上下文表示）
1. 输入特征图 $X \in \mathbb{R}^{C \times H \times W}$ 被转换为两个部分： $Q \in \mathbb{R}^{C \times HW}$ 和 $K \in \mathbb{R}^{C \times S^2}$ ，其中 K是通过自适应平均池化从 X 中聚集得到的 $S \times S$ 区域中心。
2. 通道被均匀分成 G 组，类似于多头注意力中的“头”。通过矩阵乘法计算每组的亲和力矩阵 $A_g \in \mathbb{R}^{HW \times S^2}$ ，其中 $A_{gi}$ 表示 $Q_g$ 中第 i 个token与 $K_g$ 中所有token之间的亲和力值。
Token-wise Global Context Mixing（全局上下文混合）
1. 使用另一个可学习的线性层 $W_d \in \mathbb{R}^{S^2 \times K^2}$ 对每个亲和力矩阵 $A_{g}$ 的行进行聚合，生成动态卷积核。
2. 使用softmax函数对聚合后的亲和力值进行归一化，得到 $D_g = \text{softmax}(A_g W_d) \in \mathbb{R}^{HW \times K^2}$ 。将 $D_g$ 的每一行重塑为目标卷积核形状，生成每个token位置的输入依赖卷积核。
Implementation（实现细节）
1. 在OverLoCK网络的Dynamic Block中，ContMix被定制化实现。具体来说，Q 和 K 矩阵分别使用 $Z_i$ （当前特征图）和 $P_i$ （上下文先验）的通道来计算。
2. 设置 S=7 以确保ContMix具有线性时间复杂度。同时，为了更好地提取多尺度特征，将一半的组分配给大卷积核（如5×5），另一半分配给小卷积核（如3×3）。
3. 使用Dilated RepConv层增加通道多样性

ContMix通过将全局上下文信息注入到卷积核权重中，有效地解决了传统卷积在处理大输入分辨率时的局限性。它不仅能够建模长距离依赖关系，还能保持强大的局部归纳偏差，从而在各种视觉任务中实现更好的性能。

Dilated RepConv

Dilated RepConv 是一种结合了空洞卷积（扩张卷积，Dilated Convolution）和重参数化（RepConv）技术的卷积模块，广泛应用于计算机视觉任务中，特别是在目标检测和图像分类领域。它通过在训练阶段使用多分支结构，并在推理阶段将这些结构简化为单一的等效卷积操作，从而在保持高性能的同时显著降低计算开销。

扩张卷积通过在卷积核中引入“空洞”来扩大感受野，从而在不增加额外参数的情况下增加模型的覆盖范围。例如，一个3×3的扩张卷积核可以通过设置扩张率（dilation rate）来“跳过”某些像素，从而覆盖更大的区域。这种技术特别适用于需要捕捉更大范围信息的任务，如语义分割和目标检测。

空洞卷积参考资料：空洞卷积(dilated convolution)深入详解——优点与缺点-腾讯云开发者社区-腾讯云

重参数化技术的核心思想是将训练时的多分支结构在推理时简化为单一的卷积操作。具体来说，RepConv模块通常包含以下几种路径：

一个标准的3×3卷积路径。
一个1×1卷积路径，用于增加通道数或进行特征融合。
一个恒等映射（Identity Mapping），直接传递输入数据。

在训练阶段，这些路径共同作用，提升了模型的表达能力和泛化能力。在推理阶段，这些路径被合并为一个等效的3×3卷积核，显著减少了计算量。

见：https://arxiv.org/pdf/2101.03697

ConvFFN

ConvFFN 是一种结合了卷积神经网络（CNN）和全连接前馈网络（FFN）的结构，它通过引入卷积操作来增强模型对局部特征的捕捉能力，同时保持了FFN的高效性和灵活性。

ConvFFN的设计受到Transformer中前馈网络（FFN）的启发，但与传统的FFN不同，它在全连接层之前使用了卷积操作（如3×3卷积）来提取局部特征。这种设计不仅能够捕捉到更丰富的局部信息，还能在不显著增加计算量的情况下提高模型的性能。

SE Layer

SE Layer 是一种用于深度神经网络的注意力机制模块，旨在通过显式地建模特征通道之间的相互依赖关系来增强网络的表达能力。它通过动态调整不同通道的重要性权重，使模型能够更好地关注重要的特征通道，从而提高性能。

SE Layer 的工作流程可以分为三个主要阶段：Squeeze（挤压）、Excitation（激励）和 Scale（缩放）。

Squeeze（挤压）：通过全局平均池化（Global Average Pooling）将输入特征图的每个通道压缩为一个单一的数值，从而生成一个全局描述符。这一步骤将特征图的空间维度（高度和宽度）压缩为1，保留了每个通道的全局信息。
Excitation（激励）：使用两个全连接层（或1×1卷积层）对挤压后的全局描述符进行处理，生成每个通道的重要性权重。具体来说，先通过一个降维的全连接层（或1×1卷积层）减少参数数量，然后通过ReLU激活函数增加非线性，最后通过一个升维的全连接层（或1×1卷积层）恢复到原始通道数，并通过Sigmoid函数将输出值限制在0到1之间。
Scale（缩放）：将激励阶段生成的通道权重通过逐通道乘法应用到输入特征图上，从而对每个通道的特征进行加权。这一步骤使得重要的通道特征被增强，而不重要的通道特征被抑制。

训练策略

在ImageNet-1K数据集上进行预训练时，Overview-Net和Focus-Net各自连接一个分类器头，并施加相同的分类损失，以实现表示学习。

在迁移到下游任务时，不再对Overview-Net施加辅助监督信号，因为它已经在预训练阶段学习了高层次表示。Focus-Net用于分类任务的预测，而在密集预测任务中，使用Base-Net和Focus-Net的特征构建特征金字塔。

实验结果

图像分类

实验设置：作者在ImageNet-1K数据集上进行了图像分类实验，遵循DeiT的训练设置，使用AdamW优化器训练300个epoch。为了评估不同大小的OverLoCK模型，作者分别测试了Extreme-Tiny（XT）、Tiny（T）、Small（S）和Base（B）四个变体。

OverLoCK-XT在Top-1准确率上达到了82.7%，超过了BiFormer-T（81.4%）和UniRepLKNet-N（81.6%）。
OverLoCK-T达到了84.2%的Top-1准确率，优于MogaNet-S（83.4%）和PeLK-T（82.6%）。
OverLoCK-S达到了84.8%的Top-1准确率，超过了BiFormer-B（84.3%）和UniRepLKNet-S（83.9%）。
OverLoCK-B达到了85.1%的Top-1准确率，优于MaxViT-B（84.9%）。

此外，OverLoCK在推理速度上也表现出色。例如，OverLoCK-S在单个NVIDIA L40S GPU上的吞吐量比MogaNet-B高出超过100张图像/秒，同时Top-1准确率从84.3%提升到84.8%。这些结果表明，OverLoCK在性能和效率之间取得了良好的平衡。

目标检测

实验设置：作者使用COCO 2017数据集，分别采用Mask R-CNN和Cascade Mask R-CNN框架进行目标检测和实例分割任务。所有模型均在ImageNet-1K上进行预训练，然后在COCO数据集上进行微调。

使用Mask R-CNN（1×调度）时，OverLoCK-S的APb（平均精度）达到了48.3%，超过了BiFormer-B（46.7%）和MogaNet-B（46.7%）。
使用Cascade Mask R-CNN（3×调度）时，OverLoCK-S的APb达到了53.6%，超过了PeLK-S（52.2%）和UniRepLKNet-S（53.0%）。
OverLoCK-B在3×调度下达到了53.9%的APb，优于MogaNet-L（53.3%）。

这些结果表明，OverLoCK在目标检测和实例分割任务上也具有显著的性能优势，尤其是在处理高分辨率输入时，OverLoCK能够更有效地捕捉长距离依赖关系，从而提高检测精度。

语义分割

实验设置：作者在ADE20K数据集上使用UperNet框架进行语义分割任务。所有模型均使用在ImageNet-1K上预训练的权重，并遵循相同的训练设置。

实验结果：

OverLoCK-T在mIoU（平均交并比）上达到了50.3%，超过了MogaNet-S（49.2%）和UniRepLKNet-T（48.6%）。
OverLoCK-S达到了51.3%的mIoU，优于BiFormer-B（51.0%）和UniRepLKNet-S（50.5%）。
OverLoCK-B达到了51.7%的mIoU，优于MogaNet-L（50.9%）。

这些结果表明，OverLoCK在语义分割任务上也表现出色，尤其是在处理高分辨率输入时，其长距离依赖建模能力显著提升了分割精度。

3.代码详解

环境安装

# Environments:
cuda==12.1
python==3.10
# Dependencies:
pip install torch==2.3.1 torchvision==0.18.1 --index-url https://download.pytorch.org/whl/cu121
pip install natten==0.17.1+torch230cu121 -f https://shi-labs.com/natten/wheels/
pip install timm==0.6.12
pip install mmengine==0.2.0

为了加速训练和推理，作者利用了 RepLKNet 中提出的高效大核卷积。请按照https://github.com/VITA-Group/SLaK#installation安装 depthwise_conv2d_implicit_gemm 函数。

如果在安装 natten 时遇到网络问题，请下载此软件包（natten）并在本地安装。

你可以安装以下方式准备数据集

│imagenet/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......

不同版本的权重文件如下：

Models	Input Size	FLOPs (G)	Params (M)	Top-1 (%)	Download
OverLoCK-XT	224x224	2.6	16	82.7	model
OverLoCK-T	224x224	5.5	33	84.2	model
OverLoCK-S	224x224	9.7	56	84.8	model
OverLoCK-B	224x224	16.7	95	85.1	model

模型架构

主干网络

下面是OverLoCK网络的主干架构，依次调用三个不同的特征提取方法，逐步处理输入张量 x，并传递上下文信息：

x = self.forward_pre_features(x)：对输入进行初步特征提取（Base-Net的前两个模块）；
x, ctx = self.forward_base_features(x)：提取基础特征，并生成上下文信息 ctx（Base-Net的最后一个模块和Overview-Net）；
x, ctx_cls = self.forward_sub_features(x, ctx)：基于前面的输出和上下文进一步提取子特征，并输出分类相关的上下文 ctx_cls（Focus-Net）；
最终返回处理后的特征 x 和分类上下文 ctx_cls。

class OverLoCK(nn.Module):
    def __init__(self, ):
        
        ...
        
        for i in range(depth[0]):
            self.blocks1.append(RepConvBlock(...))
        
        for i in range(depth[1]):
            self.blocks2.append(RepConvBlock(...))
            
        for i in range(depth[2]):
            self.blocks3.append(RepConvBlock(...))

        for i in range(depth[3]):
            self.blocks4.append(RepConvBlock(...))
            
        for i in range(sub_depth[0]):
            self.sub_blocks3.append(DynamicConvBlock(...))
        
        for i in range(sub_depth[1]):
            self.sub_blocks4.append(
                DynamicConvBlock(...))

        # Aux Cls Head
        if use_ds:
            self.aux_head = nn.Sequential(
                nn.BatchNorm2d(embed_dim[-1]),
                nn.AdaptiveAvgPool2d(1),
                nn.Conv2d(embed_dim[-1], num_classes, kernel_size=1) if num_classes > 0 else nn.Identity()
            )
        
        # Main Cls Head
        self.head = nn.Sequential(
            nn.Conv2d(fusion_dim, projection, kernel_size=1, bias=False),
            nn.BatchNorm2d(projection),
            nn.SiLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(projection, num_classes, kernel_size=1) if num_classes > 0 else nn.Identity()
        )
        
    ...
            
    def forward_pre_features(self, x):
        x = self.patch_embed1(x)        # Embed
        for blk in self.blocks1:        # Basic Block N1(Base)
            x = blk(x)
            
        x = self.patch_embed2(x)
        for blk in self.blocks2:        # Basic Block N2(Base)
            x = blk(x)

        return x
    
    
    def forward_base_features(self, x):
        x = self.patch_embed3(x)
        for blk in self.blocks3:        # Basic Block N3(Base)
            x = blk(x)
            
        ctx = self.patch_embed4(x)
        for blk in self.blocks4:        # Basic Block N4(Overview)
            ctx = blk(ctx)

        return (x, ctx)
    

    def forward_sub_features(self, x, ctx):
        ctx_cls = ctx
        ctx_ori = self.high_level_proj(ctx)
        ctx_up = F.interpolate(ctx_ori, size=x.shape[2:], mode='bilinear', align_corners=False)
        
        for idx, blk in enumerate(self.sub_blocks3):
            if idx == 0:
                ctx = ctx_up
            x, ctx = blk(x, ctx, ctx_up)        # Dynamic Block N5(Focus)

        x, ctx = self.patch_embedx(x, ctx)
        for idx, blk in enumerate(self.sub_blocks4):
            x, ctx = blk(x, ctx, ctx_ori)       # Dynamic Block N6(Focus)
        
        return (x, ctx_cls)

    def forward_features(self, x):
        x = self.forward_pre_features(x)        # 对输入进行初步特征提取（Base-Net）
        x, ctx = self.forward_base_features(x)      # 提取基础特征，并生成上下文信息 ctx（Overview-Net）
        x, ctx_cls = self.forward_sub_features(x, ctx)      # 提取子特征，并生成子特征信息 ctx_cls（Focus-Net）

        return (x, ctx_cls)

    def forward(self, x):
        x, ctx = self.forward_features(x)
        x = self.head(x).flatten(1)         # 分类头

        if hasattr(self, 'aux_head') and self.training:
            ctx = self.aux_head(ctx).flatten(1)     # 辅助头
            return dict(main=x, aux=ctx)
        
        return x

RepConvBlock

其中RepConvBlock如下：该模块实现了一个带有残差连接、深度卷积、膨胀重参数化和前馈网络的块结构。

class RepConvBlock(nn.Module):
    def __init__():
        ...
        
        self.dwconv = ResDWConv(dim, kernel_size=3)     # 带残差的深度卷积
    
        self.proj = nn.Sequential(
            norm_layer(dim),
            DilatedReparamBlock(dim, kernel_size=kernel_size, deploy=deploy, use_sync_bn=False, attempt_use_lk_impl=use_gemm),      # DilatedRepConv
            nn.BatchNorm2d(dim),
            SEModule(dim),          # SE
            nn.Conv2d(dim, mlp_dim, kernel_size=1),         # ConvFFN
            nn.GELU(),
            ResDWConv(mlp_dim, kernel_size=3),              # ConvFFN
            GRN(mlp_dim),
            nn.Conv2d(mlp_dim, dim, kernel_size=1),         # ConvFFN
            DropPath(drop_path) if drop_path > 0 else nn.Identity(),
        )

        self.ls = LayerScale(dim, init_value=ls_init_value) if ls_init_value is not None else nn.Identity()     # 缩放和平移
        
    def forward_features(self, x):
        x = self.dwconv(x)      # 带残差的深度卷积
        
        if self.res_scale:
            x = self.ls(x) + self.proj(x)       # 缩放和平移  +  Basic Block后半部分
        else:
            drop_path = self.proj[-1]
            x = x + drop_path(self.ls(self.proj[:-1](x)))

        return x
    
    def forward(self, x):
        if self.use_checkpoint and x.requires_grad:     # 如果启用了检查点（self.use_checkpoint为真）且输入需要梯度计算（x.requires_grad为真）
            x = checkpoint(self.forward_features, x, use_reentrant=False)       # 则使用checkpoint函数包装前向传播以节省内存
        else:
            x = self.forward_features(x)        # 直接调用self.forward_features处理输入
        
        return x

该类实现了一个带有残差连接的深度可分离卷积层。

class ResDWConv(nn.Conv2d):
    '''
    Depthwise convolution with residual connection
    '''
    def __init__(self, dim, kernel_size=3):
        super().__init__(dim, dim, kernel_size=kernel_size, padding=kernel_size//2, groups=dim)
    
    def forward(self, x):
        x = x + super().forward(x)
        return x

该代码实现了一个全局响应归一化（GRN）层，主要用于神经网络中对输入特征进行通道上的归一化和仿射变换。

class GRN(nn.Module):       # 全局响应归一化（GRN）层，主要用于神经网络中对输入特征进行通道上的归一化和仿射变换。
    def __init__(self, dim, use_bias=True):
        super().__init__()
        self.use_bias = use_bias
        self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1))
        if self.use_bias:
            self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1))

    def forward(self, x):
        Gx = torch.norm(x, p=2, dim=(-1, -2), keepdim=True)     # 计算输入 x 在空间维度上的 L2 范数 Gx
        Nx = Gx / (Gx.mean(dim=1, keepdim=True) + 1e-6)         # 对 Gx 在通道维度上进行归一化得到 Nx
        if self.use_bias:
            return (self.gamma * Nx + 1) * x + self.beta        # 使用 gamma 和 beta 对输入进行仿射变换并返回结果。
        else:
            return (self.gamma * Nx + 1) * x

该代码定义了一个 LayerScale 模块，功能是对输入张量 x 在通道维度上进行可学习的缩放和平移操作。

class LayerScale(nn.Module):        # 对输入张量 x 在通道维度上进行可学习的缩放和平移操作。
    def __init__(self, dim, init_value=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim, 1, 1, 1)*init_value, 
                                   requires_grad=True)
        self.bias = nn.Parameter(torch.zeros(dim), requires_grad=True)

    def forward(self, x):
        x = F.conv2d(x, weight=self.weight, bias=self.bias, groups=x.shape[1])
        return x

DynamicConvBlock

该函数实现了一个动态卷积块，整体结构用于高效地融合多尺度上下文信息并提升模型表达能力。

class DynamicConvBlock(nn.Module):
    def __init__(...):
        
        ...
        
        self.dwconv1 = ResDWConv(out_dim, kernel_size=3)
        self.norm1 = norm_layer(out_dim)
        
        self.fusion = nn.Sequential(
            nn.Conv2d(out_dim, out_dim, kernel_size=3, padding=1, groups=out_dim),
            nn.BatchNorm2d(out_dim),
            nn.GELU(),
            nn.Conv2d(out_dim, dim, kernel_size=1),
            GRN(dim),
        )
        
        self.weight_query = nn.Sequential(
            nn.Conv2d(dim, dim//2, kernel_size=1, bias=False),
            nn.BatchNorm2d(dim//2),
        )
         
        self.weight_key = nn.Sequential(
            nn.AdaptiveAvgPool2d(7),
            nn.Conv2d(ctx_dim, dim//2, kernel_size=1, bias=False),
            nn.BatchNorm2d(dim//2),
        )
        
        self.weight_proj = nn.Conv2d(49, kernel_size**2 + smk_size**2, kernel_size=1)
        
        self.dyconv_proj = nn.Sequential(
            nn.Conv2d(dim, dim, kernel_size=1, bias=False),
            nn.BatchNorm2d(dim),
        )
        
        self.lepe = nn.Sequential(
            DilatedReparamBlock(dim, kernel_size=kernel_size, deploy=deploy, use_sync_bn=False, attempt_use_lk_impl=use_gemm),
            nn.BatchNorm2d(dim),
        )
        
        self.se_layer = SEModule(dim)
        
        self.gate = nn.Sequential(
            nn.Conv2d(dim, dim, kernel_size=1, bias=False),
            nn.BatchNorm2d(dim),
            nn.SiLU(),
        )

        self.proj = nn.Sequential(
            nn.BatchNorm2d(dim),
            nn.Conv2d(dim, out_dim, kernel_size=1),
        )
        
        self.dwconv2 = ResDWConv(out_dim, kernel_size=3)
        self.norm2 = norm_layer(out_dim)
        
        self.mlp = nn.Sequential(
            nn.Conv2d(out_dim, mlp_dim, kernel_size=1),
            nn.GELU(),
            ResDWConv(mlp_dim, kernel_size=3),
            GRN(mlp_dim),
            nn.Conv2d(mlp_dim, out_dim, kernel_size=1),
        )
        
        self.ls1 = LayerScale(out_dim, init_value=ls_init_value) if ls_init_value is not None else nn.Identity()
        self.ls2 = LayerScale(out_dim, init_value=ls_init_value) if ls_init_value is not None else nn.Identity()
        self.drop_path = DropPath(drop_path) if drop_path > 0 else nn.Identity()
        
        self.get_rpb()


    def get_rpb(self):
        self.rpb_size1 = 2 * self.smk_size - 1
        self.rpb1 = nn.Parameter(torch.empty(self.num_heads, self.rpb_size1, self.rpb_size1))
        self.rpb_size2 = 2 * self.kernel_size - 1
        self.rpb2 = nn.Parameter(torch.empty(self.num_heads, self.rpb_size2, self.rpb_size2))
        nn.init.zeros_(self.rpb1)
        nn.init.zeros_(self.rpb2)
    
        
    @torch.no_grad()
    def generate_idx(self, kernel_size):
        rpb_size = 2 * kernel_size - 1
        idx_h = torch.arange(0, kernel_size)
        idx_w = torch.arange(0, kernel_size)
        idx_k = ((idx_h.unsqueeze(-1) * rpb_size) + idx_w).view(-1)     # 将二维卷积核展开为一维索引，用于后续特征图的动态采样。
        return (idx_h, idx_w, idx_k)
    

    def apply_rpb(self, attn, rpb, height, width, kernel_size, idx_h, idx_w, idx_k):
        """
        将相对位置偏置（RPB）应用到注意力矩阵上
        """
        num_repeat_h = torch.ones(kernel_size, dtype=torch.long)
        num_repeat_w = torch.ones(kernel_size, dtype=torch.long)
        num_repeat_h[kernel_size//2] = height - (kernel_size-1)
        num_repeat_w[kernel_size//2] = width - (kernel_size-1)
        bias_hw = (idx_h.repeat_interleave(num_repeat_h).unsqueeze(-1) * (2*kernel_size-1)) + idx_w.repeat_interleave(num_repeat_w)
        bias_idx = bias_hw.unsqueeze(-1) + idx_k
        bias_idx = bias_idx.reshape(-1, int(kernel_size**2))
        bias_idx = torch.flip(bias_idx, [0])
        rpb = torch.flatten(rpb, 1, 2)[:, bias_idx]
        rpb = rpb.reshape(1, int(self.num_heads), int(height), int(width), int(kernel_size**2))
        return attn + rpb
    

    def _forward_inner(self, x, h_x, h_r):
        input_resoltion = x.shape[2:]   
        B, C, H, W = x.shape
        B, C_h, H_h, W_h = h_x.shape
        
        if not self.is_first:
            h_x = self.x_scale(h_x) + self.h_scale(h_r)

        x_f = torch.cat([x, h_x], dim=1)        # Zi+Pi  [B, C+C_h, H, W]
        x_f = self.dwconv1(x_f)             # 深度可分离卷积
        identity = x_f
        x_f = self.norm1(x_f)
        x = self.fusion(x_f)
        gate = self.gate(x)
        lepe = self.lepe(x)

        is_pad = False
        if min(H, W) < self.kernel_size:
            is_pad = True
            if H < W:
                size = (self.kernel_size, int(self.kernel_size / H * W))
            else:
                size = (int(self.kernel_size / W * H), self.kernel_size)
            x = F.interpolate(x, size=size, mode='bilinear', align_corners=False)
            x_f = F.interpolate(x_f, size=size, mode='bilinear', align_corners=False)
            H, W = size

        # ContMix
        query, key = torch.split(x_f, split_size_or_sections=[C, C_h], dim=1)       # 将输入特征 x_f 按通道分为 query 和 key 两部分
        query = self.weight_query(query) * self.scale       # 线性变换
        key = self.weight_key(key)
        query = rearrange(query, 'b (g c) h w -> b g c (h w)', g=self.num_heads)        # 多头重排
        key = rearrange(key, 'b (g c) h w -> b g c (h w)', g=self.num_heads)
        weight = einsum(query, key, 'b g c n, b g c l -> b g n l')
        weight = rearrange(weight, 'b g n l -> b l g n').contiguous()       # 相似度计算
        weight = self.weight_proj(weight)           # 投影与重构
        weight = rearrange(weight, 'b l g (h w) -> b g h w l', h=H, w=W)

        attn1, attn2 = torch.split(weight, split_size_or_sections=[self.smk_size**2, self.kernel_size**2], dim=-1)      # 权重拆分
        rpb1_idx = self.generate_idx(self.smk_size)     # 生成用于相对位置偏置（RPB）的索引
        rpb2_idx = self.generate_idx(self.kernel_size)
        attn1 = self.apply_rpb(attn1, self.rpb1, H, W, self.smk_size, *rpb1_idx)        # 应用RPB
        attn2 = self.apply_rpb(attn2, self.rpb2, H, W, self.kernel_size, *rpb2_idx)
        attn1 = torch.softmax(attn1, dim=-1)        # Softmax归一化
        attn2 = torch.softmax(attn2, dim=-1)
        value = rearrange(x, 'b (m g c) h w -> m b g h w c', m=2, g=self.num_heads)     # 值向量重排

        x1 = na2d_av(attn1, value[0], kernel_size=self.smk_size)        # 动态卷积块中的两个非对称注意力加权特征提取
        x2 = na2d_av(attn2, value[1], kernel_size=self.kernel_size)

        x = torch.cat([x1, x2], dim=1)
        x = rearrange(x, 'b g h w c -> b (g c) h w', h=H, w=W)

        if is_pad:
            x = F.adaptive_avg_pool2d(x, input_resoltion)

        x = self.dyconv_proj(x)

        x = x + lepe
        x = self.se_layer(x)        # SE

        x = gate * x
        x = self.proj(x)

        if self.res_scale:
            x = self.ls1(identity) + self.drop_path(x)
        else:
            x = identity + self.drop_path(self.ls1(x))
        
        x = self.dwconv2(x)
         
        if self.res_scale:
            x = self.ls2(x) + self.drop_path(self.mlp(self.norm2(x)))
        else:
            x = x + self.drop_path(self.ls2(self.mlp(self.norm2(x))))

        if self.is_last:
            return (x, None)
        else:
            l_x, h_x = torch.split(x, split_size_or_sections=[C, C_h], dim=1)
            return (l_x, h_x)
    
    def forward(self, x, h_x, h_r):
        if self.use_checkpoint and x.requires_grad:
            x = checkpoint(self._forward_inner, x, h_x, h_r, use_reentrant=False)
        else:
            x = self._forward_inner(x, h_x, h_r)
        return x

4.总结

OverLoCK的设计包含三个协同工作的子网络：Base-Net负责编码低层次和中层次的特征；Overview-Net通过粗粒度的全局上下文建模生成动态的自上而下的注意力；Focus-Net则在自上而下的注意力引导下进行更细致的特征感知。为了充分发挥自上而下注意力的潜力，作者还提出了一种新颖的上下文混合动态卷积（ContMix），它能够在保持局部归纳偏差的同时有效建模长距离依赖关系，即使在输入分辨率增加时也能保持性能。

实验结果表明，OverLoCK在多个视觉任务上均取得了显著的性能提升。例如，在ImageNet-1K数据集上，OverLoCK-T模型达到了84.2%的Top-1准确率，显著超过了ConvNeXt-B等现有方法，同时仅使用了大约三分之一的计算量和参数。在目标检测任务中，OverLoCK-S模型在COCO数据集上的APb指标比MogaNet-B高出1%。在语义分割任务中，OverLoCK-T模型在ADE20K数据集上的mIoU指标比UniRepLKNet-T高出1.7%。此外，OverLoCK还能够生成更大的有效感受野，并在特征响应上表现出更强的局部归纳偏差。代码已公开，供研究者们进一步探索和应用。

如果你觉得这篇文章对你有帮助，或者你对这个话题感兴趣，不妨点个赞、收藏一下，或者关注我哦！这样不仅能支持我继续创作，还能让你第一时间获取更多优质内容！感谢你的支持！💖

OverLoCK：先概览，再聚焦。CVPR2025全新主干网络