DCN可形变卷积实现1：Python实现

DCN可形变卷积实现1：Python实现

我们会先用纯 Python 实现一个 Pytorch 版本的 DCN ，然后实现其 C++/CUDA 版本。

本文主要关注 DCN 可形变卷积的代码实现，不会过多的介绍其思想，如有兴趣，请参考论文原文：

Deformable Convolutional Networks

Deformable ConvNets v2: More Deformable, Better Results

DCN简介

考虑到传统卷积必须是方方正正的 $k\times k$ 的卷积核：
$$\mathbf{y}({\mathbf{p}}_0)=\sum _{{\mathbf{p}}_n\in \mathcal{R}}\mathbf{w}({\mathbf{p}}_n)\cdot \mathbf{x}({\mathbf{p}}_0+{\mathbf{p}}_n)$$
作者认为这个感受野太规则，无法很好地捕捉特殊形状的特征，因此在其基础上加了偏置：
$$\mathbf{y}({\mathbf{p}}_0)=\sum _{{\mathbf{p}}_n\in \mathcal{R}}\mathbf{w}({\mathbf{p}}_n)\cdot \mathbf{x}({\mathbf{p}}_0+{\mathbf{p}}_n+\Delta {\mathbf{p}}_n)$$
使得模型能够根据输入计算偏移量，自己选择对哪些位置进行卷积计算，而不用必须是正方形的样子。

如上图所示，传统的卷积输入只能是图 (a) 中的九个绿点，而在加上偏移量之后，皆可以四处飞，比如飞到图 (bcd) 中蓝点的位置。

而 DCNv2 则在此基础上又为每个位置乘了一个可学习的权重：
$$\mathbf{y}({\mathbf{p}}_0)=\sum _{{\mathbf{p}}_n\in \mathcal{R}}\mathbf{w}({\mathbf{p}}_n)\cdot \mathbf{x}({\mathbf{p}}_0+{\mathbf{p}}_n+\Delta {\mathbf{p}}_n)\cdot \Delta {\mathbf{m}}_n$$

由于网络学习出的偏移量通常是小数，因此下面会用到双线性插值（下面会有图示），这里先把原文中的公式给出来：
$$\mathbf{x}(\mathbf{p})=\sum _{\mathbf{q}}G(\mathbf{q},\mathbf{p})\cdot \mathbf{x}(\mathbf{q})$$

这里 $\mathbf{p}=({\mathbf{p}}_0+{\mathbf{p}}_n+\Delta {\mathbf{p}}_n)$ 表示任意位置（可以是小数）坐标，而 $\mathbf{q}$ 是枚举特征图 $\mathbf{x}$ 中所有整数空间位置，$G(\cdot ,\cdot )$ 就是双线性插值，注意这里的 $G$ 是两个维度（x,y）的，拆分为两个单维度的话，就是：

$$G(\mathbf{q},\mathbf{p})=g(q_x,p_x)\cdot g(q_y,p_y)$$

其中 $g(a,b)=max(0,1-|a-b|)$ 。

给出公式一方面是让读者了解具体算法，更重要的一点是我们参考的 DCN 的 Pytorch 实现代码中变量的命名是与原文公式对应的，因此公式列在这里方便读者下面看代码的时候可以回头看一下各个变量对应的是算法公式中的哪一项。

纯Python实现

我们先来看一下Pytorch版本的实现，来更好地理解 DCN 可形变卷积的做法，然后用 C++/CUDA 实现高性能版本。本文参考的 Python 实现是：https://github.com/4uiiurz1/pytorch-deform-conv-v2/blob/master/deform_conv_v2.py 。

本小节参考博文：deformable convolution可变形卷积（4uiiurz1-pytorch版）源码分析

_init_

def __init__(self, inc, outc, kernel_size=3, padding=1, stride=1, bias=None, modulation=False):
    """
    Args:
        modulation (bool, optional): If True, Modulated Defomable Convolution (Deformable ConvNets v2).
    """
    super(DeformConv2d, self).__init__()
    self.kernel_size = kernel_size
    self.padding = padding
    self.stride = stride
    self.zero_padding = nn.ZeroPad2d(padding)
    self.conv = nn.Conv2d(inc, outc, kernel_size=kernel_size, stride=kernel_size, bias=bias)

    self.p_conv = nn.Conv2d(inc, 2*kernel_size*kernel_size, kernel_size=3, padding=1, stride=stride)
    nn.init.constant_(self.p_conv.weight, 0)
    self.p_conv.register_backward_hook(self._set_lr)

    self.modulation = modulation
    if modulation:
        self.m_conv = nn.Conv2d(inc, kernel_size*kernel_size, kernel_size=3, padding=1, stride=stride)
        nn.init.constant_(self.m_conv.weight, 0)
        self.m_conv.register_backward_hook(self._set_lr)

这里重点关注 self.p_conv 和 self.m_conv ，是这两个卷积完成了对偏移量 offset 的学习，而 self.conv 是确在定偏移后的位置之后，最终进行计算的卷积。

（关于这里的 modulation 参数，如注释所言，如果为 True ，就是一个模块化的 DCN，即 DCNv2。）

具体来看这三个卷积及其参数：

• self.conv：这是负责进行最终计算的卷积。可形变卷积 DCN 虽然进行了形变，但是这是卷积输入中空间像素的位置有了偏移，而输入输出的尺寸还是不变的，因此，输入卷积的位置确定之后，最终负责完成卷积计算的 self.conv 的各个参数（输入输出通道数inc, outc、卷积核大小kernel_size、步长stride、填充padding等）就是我们整个 DCN 的对应参数参数。

• self.p_conv：该卷积操作负责计算偏移量。在卷积中，共有 kernel_size * kernel_size 个位置的像素需要参与计算，因此我们要计算出他们的偏移量，而每个位置都有宽、高两个方向的偏移量，故该卷积输出的通道数是 2 * kernel_size * kernel_size ，其他参数保持一致。

• self.m_conv：该卷积操作负责计算卷积核每个位置的权重。其输出通道数为位置数，即 kernel_size * kernel_size ，其他参数保持一致，注意这个加权的想法是 DCNv2 中的。

forward

看过 __init__ 函数之后，我们可以来看 forward 函数：

def forward(self, x):
    offset = self.p_conv(x)
    if self.modulation:
        m = torch.sigmoid(self.m_conv(x))

    dtype = offset.data.type()
    ks = self.kernel_size
    N = offset.size(1) // 2

    if self.padding:
        x = self.zero_padding(x)

    # (b, 2N, h, w)
    p = self._get_p(offset, dtype)

    # (b, h, w, 2N)
    p = p.contiguous().permute(0, 2, 3, 1)
    q_lt = p.detach().floor()
    q_rb = q_lt + 1

    q_lt = torch.cat([torch.clamp(q_lt[..., :N], 0, x.size(2)-1), torch.clamp(q_lt[..., N:], 0, x.size(3)-1)], dim=-1).long()
    q_rb = torch.cat([torch.clamp(q_rb[..., :N], 0, x.size(2)-1), torch.clamp(q_rb[..., N:], 0, x.size(3)-1)], dim=-1).long()
    q_lb = torch.cat([q_lt[..., :N], q_rb[..., N:]], dim=-1)
    q_rt = torch.cat([q_rb[..., :N], q_lt[..., N:]], dim=-1)

    # clip p
    p = torch.cat([torch.clamp(p[..., :N], 0, x.size(2)-1), torch.clamp(p[..., N:], 0, x.size(3)-1)], dim=-1)

    # bilinear kernel (b, h, w, N)
    g_lt = (1 + (q_lt[..., :N].type_as(p) - p[..., :N])) * (1 + (q_lt[..., N:].type_as(p) - p[..., N:]))
    g_rb = (1 - (q_rb[..., :N].type_as(p) - p[..., :N])) * (1 - (q_rb[..., N:].type_as(p) - p[..., N:]))
    g_lb = (1 + (q_lb[..., :N].type_as(p) - p[..., :N])) * (1 - (q_lb[..., N:].type_as(p) - p[..., N:]))
    g_rt = (1 - (q_rt[..., :N].type_as(p) - p[..., :N])) * (1 + (q_rt[..., N:].type_as(p) - p[..., N:]))

    # (b, c, h, w, N)
    x_q_lt = self._get_x_q(x, q_lt, N)
    x_q_rb = self._get_x_q(x, q_rb, N)
    x_q_lb = self._get_x_q(x, q_lb, N)
    x_q_rt = self._get_x_q(x, q_rt, N)

    # (b, c, h, w, N)
    x_offset = g_lt.unsqueeze(dim=1) * x_q_lt + \
               g_rb.unsqueeze(dim=1) * x_q_rb + \
               g_lb.unsqueeze(dim=1) * x_q_lb + \
               g_rt.unsqueeze(dim=1) * x_q_rt

    # modulation
    if self.modulation:
        m = m.contiguous().permute(0, 2, 3, 1)
        m = m.unsqueeze(dim=1)
        m = torch.cat([m for _ in range(x_offset.size(1))], dim=1)
        x_offset *= m

    x_offset = self._reshape_x_offset(x_offset, ks)
    out = self.conv(x_offset)

    return out

这里的 N 是 offset 的通道数除以2，就是卷积要处理的位置的个数（即 kernal_size * kernel_size）。

整个 forward 函数的流程：

1. 首先通过上面介绍的 p_conv 和 v_conv 计算出偏移量 offset 和加权的权重m（如果有）。

2. 比较关键的是这里的 self._get_p 函数，该函数通过上面计算出的 offset，去得到输入到卷积的具体位置，即公式中的：
   $${\mathbf{p}}_0+{\mathbf{p}}_n+\Delta {\mathbf{p}}_n$$
   关于这个函数，我们会在下一小节详细介绍。由于我们现在先过整个流程，只需要知道该函数通过 p_conv 卷积计算出的 offset，得到了要输入最终卷积的位置 p。p 是一个形状为 $(bs,2\ast N,h,w)$ 的张量。

3. 拿到 p 之后的问题是我们得到的肯定是一个浮点类型，即小数，但是像素的坐标肯定是整型，所以，这里我们需要做一个双线性插值。双线性插值的思想也很直接，就是将某个浮点坐标的左上、左下、右上、右下四个位置的像素值按照与该点的距离计算加权和，作为该点处的像素值。可参考下图，也可参考博客图像预处理之warpaffine与双线性插值及其高性能实现，后半部分有对双线性插值的讲解与 Python 实现。

   这里的 lt, rb, lb, rt 分别代表左上，右下，左下，右上。

4. 现在我们通过双线性插值拿到了每个位置的坐标，下一步就是根据坐标去取到对应位置的像素值，这在代码中由 self._get_x_q 实现，会在下面的小节介绍。

5. 这个时候如果有权重的话，要计算出 m，乘到 x_offset 上。

6. 这时得到的 x_offset 的形状是 $b,c,h,w,N$，而我们要的形状肯定是 $b,c,h,w$，因此这里还有一个 reshape 的操作，由 self._reshape_x_offset 实现。

7. 至此，我们终于得到了想要的 x_offset，接下来就将它送入 self.conv 进行卷积计算并返回结果即可。

_get_p、_get_p_0、_get_p_n

先贴一下代码：

def _get_p(self, offset, dtype):
    N, h, w = offset.size(1)//2, offset.size(2), offset.size(3)

    # (1, 2N, 1, 1)
    p_n = self._get_p_n(N, dtype)
    # (1, 2N, h, w)
    p_0 = self._get_p_0(h, w, N, dtype)
    p = p_0 + p_n + offset
    return p

def _get_p_n(self, N, dtype):
    p_n_x, p_n_y = torch.meshgrid(
        torch.arange(-(self.kernel_size-1)//2, (self.kernel_size-1)//2+1),
        torch.arange(-(self.kernel_size-1)//2, (self.kernel_size-1)//2+1))
    # (2N, 1)
    p_n = torch.cat([torch.flatten(p_n_x), torch.flatten(p_n_y)], 0)
    p_n = p_n.view(1, 2*N, 1, 1).type(dtype)

    return p_n

def _get_p_0(self, h, w, N, dtype):
    p_0_x, p_0_y = torch.meshgrid(
        torch.arange(1, h*self.stride+1, self.stride),
        torch.arange(1, w*self.stride+1, self.stride))
    p_0_x = torch.flatten(p_0_x).view(1, 1, h, w).repeat(1, N, 1, 1)
    p_0_y = torch.flatten(p_0_y).view(1, 1, h, w).repeat(1, N, 1, 1)
    p_0 = torch.cat([p_0_x, p_0_y], 1).type(dtype)

    return p_0

我们来看一下如何将 offset 传入 self._get_p 获得最终的 p，该函数会分别调用 self._get_p_0 和 self._get_p_n 来分别获得 p_0 和 p_n，分别是卷积核的中心坐标和相对坐标，对应到公式中的 ${\mathbf{p}}_0,{\mathbf{p}}_n$：
$$\mathbf{y}({\mathbf{p}}_0)=\sum _{{\mathbf{p}}_n\in \mathcal{R}}\mathbf{w}({\mathbf{p}}_n)\cdot \mathbf{x}({\mathbf{p}}_0+{\mathbf{p}}_n+\Delta {\mathbf{p}}_n)$$

关于 p_0 和 p_n 具体是什么东西其实很好理解，画个小图就明白了，以 kernel_size = 3 的卷积为例，中心位置在全图中的坐标就是 p_0，中心位置的相对坐标就是 p_n=(0,0)，左上角的 p_n=(-1,-1)，右下角的 p_n=(1,1) 其他位置以此类推。常规的卷积就只有 ${\mathbf{p}}_n+{\mathbf{p}}_0$ ，输入就是只能在上图中的九个格子中，而 DCN 加入 $\Delta {\mathbf{p}}_n$ 之后，就可以四处飞啦。但是四处飞，也是要在 ${\mathbf{p}}_n+{\mathbf{p}}_0$ 的基础上再加上偏移量来计算具体的位置。所以我们先要获得 p_0 和 p_n。

当然，p_0 和 p_n 都是固定的、不需要学习的、而且是很规则的，因此获取他们只需要根据 kernel_size 和位置 h, w （仅 p_0 需要）来计算就好了。这里代码实现中就是用 torch.arange 和 torch.meshgrid 将想要的 p_0 和 p_n，计算出来。

然后 p = p_0 + p_n + offset（对应公式），得到尺寸为 $(bs,2\ast N,h,w)$ 的 p。

_get_x_q

_get_x_q 函数是根据计算出的位置坐标，得到该位置的像素值。

再提醒一下，我们参考的 DCN 的 Pytorch 实现代码中变量的命名是与原文公式对应的，如果有变量含义不明确的，可以回上面看看公式，对应代码变量名理解。

def _get_x_q(self, x, q, N):
    b, h, w, _ = q.size()
    padded_w = x.size(3)
    c = x.size(1)
    # (b, c, h*w)
    x = x.contiguous().view(b, c, -1)

    # (b, h, w, N)
    index = q[..., :N]*padded_w + q[..., N:]  # offset_x*w + offset_y
    # (b, c, h*w*N)
    index = index.contiguous().unsqueeze(dim=1).expand(-1, c, -1, -1, -1).contiguous().view(b, c, -1)

    x_offset = x.gather(dim=-1, index=index).contiguous().view(b, c, h, w, N)

    return x_offset

_reshape_x_offset

我们在取完像素值之后得到的 x_offset 的形状是 $b,c,h,w,N$，而我们要的形状肯定是 $b,c,h,w$，因此这里还有一个 reshape 的操作，就是这里的 self._reshape_x_offset ：

@staticmethod
def _reshape_x_offset(x_offset, ks):
    b, c, h, w, N = x_offset.size()
    x_offset = torch.cat([x_offset[..., s:s+ks].contiguous().view(b, c, h, w*ks) for s in range(0, N, ks)], dim=-1)
    x_offset = x_offset.contiguous().view(b, c, h*ks, w*ks)

    return x_offset

小结

至此，我们已经使用 Pytorch 实现了纯 Python 的 DCN 卷积结构，但是，如此实现由于不是原生的 C++/CUDA 算子，而且最后的 reshape 操作虽然比较巧妙，但其实空间冗余比较大，和原文作者的 cuda 版本内存占用量差了10几倍。这个是因为在 im2col 上直接操作可以去掉很冗余。下面一篇我们会再介绍一个 C++/CUDA 实现的 DCN。

Ref

• deformable convolution可变形卷积（4uiiurz1-pytorch版）源码分析
• 图像预处理之warpaffine与双线性插值及其高性能实现