PyTorch實現深度可分離卷積（以MobileNet爲例）

介紹深度可分離卷積之前首先要先介紹一下分組卷積。

分組卷積

參考鏈接[1]已經將分組卷積介紹的比較詳細了，這裏就不再贅述。原理可以參考一下參考鏈接。但是這篇文章對於分組卷積的具體代碼講解，感覺不太好。這裏對於分組卷積着重介紹一下代碼講解。

直接調用torch.nn.Conv2d()就可以實現分組卷積。在Conv2d裏面有一個參數叫groups，這個參數是實現這一功能的關鍵。下面詳細解介紹一下這個參數。

groups

通過groups來實現分組卷積的操作。groups默認爲1，也就是說將輸入分爲一組，此時是常規卷積。groups數值爲幾表示的就是將輸入通道分爲幾組。當groups=in_channels的時候，表示的就是將輸入的每一個通道都作爲一組，然後分別對其進行卷積，輸出通道數爲k，最後再將每組輸出串聯，最後通道數爲in_channels*k

爲了詳細說明這一點，我們來舉幾個具體例子。爲了接下來方便計算，這裏bias=False，另外用到了torchsummary模塊，可以通過pip安裝。

import torch
from torchsummary import summary
import torch.nn as nn

class CSDN_Tem(nn.Module):
    def __init__(self, in_ch, out_ch, groups):
        super(CSDN_Tem, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_ch,
            out_channels=out_ch,
            kernel_size=3,
            stride=1,
            padding=1,
            groups=groups,
            bias=False
        )

    def forward(self, input):
        out = self.conv(input)
        return out

測試時，使用的是64*64*3的輸入，你可以想象爲一個64*64大小的圖片輸入。一個3*3大小的卷積核處理，最後輸出的shape是64*64*30。

我們先是group爲1的常規卷積操作。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 30, 1).to(device)
print(summary(conv,  input_size=(3, 64, 64)))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 30, 64, 64]             810
================================================================
Total params: 810
Trainable params: 810
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.94
Params size (MB): 0.00
Estimated Total Size (MB): 0.99
----------------------------------------------------------------

一共有3*3(kernel size)*3(輸入圖片的shape)*30(輸出圖片的shape)=810個參數。

下面是groups=in_channels。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 30, 3).to(device)
print(summary(conv,  input_size=(3, 64, 64)))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 30, 64, 64]             270
================================================================
Total params: 270
Trainable params: 270
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.94
Params size (MB): 0.00
Estimated Total Size (MB): 0.99
----------------------------------------------------------------

可以發現，參數量明顯的減少了。根據參考[1]的原理，過去直接卷積，卷積核大小是3*3*3，要對輸入圖片整體映射輸出30維。所以最後算的是810個參數。但是當分組爲3的時候，實際上是對輸入圖片的三個通道進行傘分組，最後輸出也是三分組。也就是對於30維的輸出而言，實際上是輸入的一個通道對應的10維輸出。這樣的話，實際我們需要3*3*1的卷積核就可以。那麼這樣參數量就變成了3*3*1*(10+10+10)=270，同樣的輸出維度，參數卻整整減小了3倍。

逐點卷積

逐點卷積就是用1*1的卷積核來對分組卷積之後的卷積網絡來進行卷積最後輸出成原來的樣子。可以看參考資料[2]有比較詳細的介紹。

深度可分離卷積

深度可分離卷積分成depthwise convolution和pointwise convolution。其中深度卷積是不只是要求group和input_channel一致。還要求input_channel和output_channel一致。也就是說對每一個通道都要單獨進行卷積操作。我們舉例來看一下：

需要特別注意的是圖片的長和寬要始終保持一致才行。

class CSDN_Tem(nn.Module):
    def __init__(self, in_ch, out_ch, kernel_size, padding, groups):
        super(CSDN_Tem, self).__init__()
        self.conv = nn.Conv2d(
            in_channels=in_ch,
            out_channels=out_ch,
            kernel_size=kernel_size,
            stride=1,
            padding=padding,
            groups=groups,
            bias=False
        )

    def forward(self, input):
        out = self.conv(input)
        return out

常規卷積是這樣的：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 30, 3, 1, 1).to(device)
print(summary(conv,  input_size=(3, 64, 64)))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 30, 64, 64]             810
================================================================
Total params: 810
Trainable params: 810
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.94
Params size (MB): 0.00
Estimated Total Size (MB): 0.99
----------------------------------------------------------------

深度可分離卷積分成兩步，第一步先是depthwise convolution：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 3, 3, padding=1, groups=3).to(device)
print(summary(conv,  input_size=(3, 64, 64)))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1            [-1, 3, 64, 64]              27
================================================================
Total params: 27
Trainable params: 27
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.09
Params size (MB): 0.00
Estimated Total Size (MB): 0.14
----------------------------------------------------------------

之後是pointwise convolution:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
conv = CSDN_Tem(3, 30, kernel_size=1, padding=0, groups=1).to(device)
print(summary(conv,  input_size=(3, 64, 64)))

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1           [-1, 30, 64, 64]              90
================================================================
Total params: 90
Trainable params: 90
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.05
Forward/backward pass size (MB): 0.94
Params size (MB): 0.00
Estimated Total Size (MB): 0.98
----------------------------------------------------------------

達到了同樣的輸出效果。但是參數量卻基本減少了8到9倍。

參考

[1] 深度可分離卷積(Depthwise Separable Convolution)和分組卷積(Group Convolution)的理解，相互關係及PyTorch實現
[2] 卷積神經網絡中的Separable Convolution