DRConv：动态区域感知卷积详解和复现

关键词：CNN、DRCNN、动态区域感知

前言

  本文将介绍一种新颖的卷积结构——Dynamic Region-Aware Convolution (动态区域感知卷积)，该结构采用了一种自适应的机制，能够在不同的感受野内自适应地调整卷积核的大小和形状，从而更好地捕捉输入中不同区域的特征信息。相比于传统的卷积结构，Dynamic Region-Aware Convolution 在各种视觉任务中都表现出了更好的性能。在本文中，我们将详细介绍该结构的设计思想和原理，并复现DRConv的结构供大家参考实验。

设计思想

  DRConv操作的输入包括三个部分：特征图、目标边界框和先验框。其中，目标边界框用于指示需要关注的区域，先验框则提供了卷积核形状和大小的参考。其核心是动态区域感知卷积，该模块可以自适应调整卷积核大小和形状，从而在感受野变化的同时保持卷积核的有效性。

流程步骤

DRConv的实现过程如下：

1. 首先根据输入特征图大小和步长计算每个卷积窗口的中心点坐标，然后以这些中心点为中心，构造多个不同大小和形状的感受野。
2. 对于每个感受野，根据其中心点坐标和感受野大小，计算相应的卷积核大小和形状，并使用这些卷积核对输入特征图进行卷积操作。
3. 最后，将所有感受野的卷积结果进行融合，得到最终的输出特征图。

优点&贡献

  DRConv 操作的主要优点是可以自适应地调整卷积核的大小和形状，以适应不同大小和形状的物体，从而提高了目标检测的精度和鲁棒性。同时，DRAC 操作还可以与现有的卷积神经网络结构无缝集成，具有较好的通用性和可拓展性。

  DRConv的核心贡献有如下三个方面：

1. 提出了一种新的动态区域感知卷积，它不仅具有强大的语义表示能力，而且完美地保持了方差特性。
2. 专门设计了可学习引导Mask的反向传播过程，以便根据反向传播的总体任务损失的梯度来确定和更新区域共享模式，这意味着本文的方法可以以端到端的方式进行优化。
3. DRConv可以通过简单地替换标准卷积而在图像分类、人脸识别、检测和分割任务上获得优异的性能，而不会增加太多计算成本。

代码：

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

from torch.autograd import Function


class asign_index(torch.autograd.Function):
    @staticmethod
    def forward(ctx, kernel, guide_feature):
        ctx.save_for_backward(kernel, guide_feature)
        guide_mask = torch.zeros_like(guide_feature).scatter_(1, guide_feature.argmax(dim=1, keepdim=True),
                                                              1).unsqueeze(2)  # B x 3 x 1 x 25 x 25
        return torch.sum(kernel * guide_mask, dim=1)

    @staticmethod
    def backward(ctx, grad_output):
        kernel, guide_feature = ctx.saved_tensors
        guide_mask = torch.zeros_like(guide_feature).scatter_(1, guide_feature.argmax(dim=1, keepdim=True),
                                                              1).unsqueeze(2)  # B x 3 x 1 x 25 x 25
        grad_kernel = grad_output.clone().unsqueeze(1) * guide_mask  # B x 3 x 256 x 25 x 25
        grad_guide = grad_output.clone().unsqueeze(1) * kernel  # B x 3 x 256 x 25 x 25
        grad_guide = grad_guide.sum(dim=2)  # B x 3 x 25 x 25
        softmax = F.softmax(guide_feature, 1)  # B x 3 x 25 x 25
        grad_guide = softmax * (grad_guide - (softmax * grad_guide).sum(dim=1, keepdim=True))  # B x 3 x 25 x 25
        return grad_kernel, grad_guide


def xcorr_slow(x, kernel, kwargs):
    """for loop to calculate cross correlation
    """
    batch = x.size()[0]
    out = []
    for i in range(batch):
        px = x[i]
        pk = kernel[i]
        px = px.view(1, px.size()[0], px.size()[1], px.size()[2])
        pk = pk.view(-1, px.size()[1], pk.size()[1], pk.size()[2])
        po = F.conv2d(px, pk, **kwargs)
        out.append(po)
    out = torch.cat(out, 0)
    return out


def xcorr_fast(x, kernel, kwargs):
    """group conv2d to calculate cross correlation
    """
    batch = kernel.size()[0]
    pk = kernel.view(-1, x.size()[1], kernel.size()[2], kernel.size()[3])
    px = x.view(1, -1, x.size()[2], x.size()[3])
    po = F.conv2d(px, pk, **kwargs, groups=batch)
    po = po.view(batch, -1, po.size()[2], po.size()[3])
    return po


class Corr(Function):
    @staticmethod
    def symbolic(g, x, kernel, groups):
        return g.op("Corr", x, kernel, groups_i=groups)

    @staticmethod
    def forward(self, x, kernel, groups, kwargs):
        """group conv2d to calculate cross correlation
        """
        batch = x.size(0)
        channel = x.size(1)
        x = x.view(1, -1, x.size(2), x.size(3))
        kernel = kernel.view(-1, channel // groups, kernel.size(2), kernel.size(3))
        out = F.conv2d(x, kernel, **kwargs, groups=groups * batch)
        out = out.view(batch, -1, out.size(2), out.size(3))
        return out


class Correlation(nn.Module):
    use_slow = True

    def __init__(self, use_slow=None):
        super(Correlation, self).__init__()
        if use_slow is not None:
            self.use_slow = use_slow
        else:
            self.use_slow = Correlation.use_slow

    def extra_repr(self):
        if self.use_slow: return "xcorr_slow"
        return "xcorr_fast"

    def forward(self, x, kernel, **kwargs):
        if self.training:
            if self.use_slow:
                return xcorr_slow(x, kernel, kwargs)
            else:
                return xcorr_fast(x, kernel, kwargs)
        else:
            return Corr.apply(x, kernel, 1, kwargs)


class DRConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, region_num=8, **kwargs):
        super(DRConv2d, self).__init__()
        self.region_num = region_num

        self.conv_kernel = nn.Sequential(
            nn.AdaptiveAvgPool2d((kernel_size, kernel_size)),
            nn.Conv2d(in_channels, region_num * region_num, kernel_size=1),
            nn.Sigmoid(),
            nn.Conv2d(region_num * region_num, region_num * in_channels * out_channels, kernel_size=1,
                      groups=region_num)
        )
        self.conv_guide = nn.Conv2d(in_channels, region_num, kernel_size=kernel_size, **kwargs)

        self.corr = Correlation(use_slow=False)
        self.kwargs = kwargs
        self.asign_index = asign_index.apply

    def forward(self, input):
        kernel = self.conv_kernel(input)
        kernel = kernel.view(kernel.size(0), -1, kernel.size(2), kernel.size(3))  # B x (r*in*out) x W X H
        output = self.corr(input, kernel, **self.kwargs)  # B x (r*out) x W x H
        output = output.view(output.size(0), self.region_num, -1, output.size(2), output.size(3))  # B x r x out x W x H
        guide_feature = self.conv_guide(input)
        output = self.asign_index(output, guide_feature)
        return output


if __name__ == "__main__":
    x1 = torch.zeros(1, 3, 640, 640)
    conv = DRConv2d(in_channels=3, out_channels=64, kernel_size=1)
    y = conv(x1)
    print(y.shape)

总结

  该方法的实现相对简单，可以通过添加一些卷积和池化层来完成，同时也可以与其他深度学习技术相结合，如目标检测、图像分类等。希望这篇博客可以为读者提供一个清晰的介绍，帮助他们理解动态区域感知卷积的工作原理和实现方法，并为他们深入研究该领域提供一些启示。

本文正在参加「金石计划」