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)?\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)?\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)?\mathbf{x}({\mathbf{p}}_0+{\mathbf{p}}_n+\Delta {\mathbf{p}}_n)?\Delta {\mathbf{m}}_n$$

由于網絡學習出的偏移量通常是小數，因此下面會用到雙線性插值（下面會有圖示），這里先把原文中的公式給出來：
$$\mathbf{x}(\mathbf{p})=\sum _{\mathbf{q}}G(\mathbf{q},\mathbf{p})?\mathbf{x}(\mathbf{q})$$

這里 $\mathbf{p}=({\mathbf{p}}_0+{\mathbf{p}}_n+\Delta {\mathbf{p}}_n)$ 表示任意位置（可以是小數）坐標，而 $\mathbf{q}$ 是枚舉特征圖 $\mathbf{x}$ 中所有整數空間位置，$G(?,?)$ 就是雙線性插值，注意這里的 $G$ 是兩個維度（x,y）的，拆分為兩個單維度的話，就是：

$$G(\mathbf{q},\mathbf{p})=g(q_x,p_x)?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_sizeself.padding = paddingself.stride = strideself.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 = modulationif 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_sizeN = offset.size(1) // 2if 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 + 1q_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 pp = 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# modulationif 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 *= mx_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?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 + offsetreturn pdef _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_ndef _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)?\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?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與雙線性插值及其高性能實現