SK-Net (Selective Kernel Networks) 架構介紹

Chinese Deep Learning

written by LiaoWC on 2021-06-11

過往做卷積（convolution）通常 kernel size 是固定的，而 SKNet[1] 的主要改進就是讓神經網路模型使用數種 kernel size 進行卷積並且能學習到如何選擇 kernel size（所以才叫 selective kernel）。SKNet 先使用多個不同 kernel size 的卷積進行運算，再用類似 SENet 的方式進行整合，讓神經網路不僅能夠學習到不同通道（channel）間的關係，還能學習到不同 kernel size 之間的關係。簡單來說可以想成：分成多個不同 kernel size 進行卷積再選擇性整合起來。至於詳細做法為何？讓我們繼續看下去。

SKNet 的架構

分支的數量可以調整，為方便討論，以下都用兩個分支的例子講解，多個分支的概念也是一樣的。

sknet-1.png

1. Split

sknet-2.png

分成數個平行的分支，每個分支皆為卷積層但各自的 kernel size 不同。以上圖為例，分成兩種 kernel size：3x3 和 5x5。實作上，5x5 不一定會真的有 25 個參數，可能用 3x3 且 dilation 為 2 的卷積層代替。Dilation 為 1 和為 2 的比較如下方兩張圖：

https://github.com/vdumoulin/conv_arithmetic/raw/master/gif/no_padding_no_strides.gif

一般的卷積(no padding, dilation = 1)

https://github.com/vdumoulin/conv_arithmetic/raw/master/gif/dilation.gif

Dilation = 2 範例

以上兩圖的來源：https://github.com/vdumoulin/conv_arithmetic [2]

所有分支用不同的 kernel size 對輸入做 grouped/depthwise convolution + Batch Normalization + ReLU（圖中兩個頭上有東西的 F），兩個分支就會產生兩個輸出，圖裡表示為兩個頭上有標記的 U。

2. Fuse

sknet-3.png

  • 把頭上有不同符號的那些的 U 各對應的元素加起來，得到 $\bold U$。
  • 對 $\bold U$ 做全域平均池化（global average pooling）得到 $\bold s$。
  • $\bold s$ 經過全連接層（fully-connected layer）+ Batch Normalization + ReLU 得到 $\bold z$。$\bold z$ 的維度為 $d=max(C/r,L)$，$r$ 是一個控制用的比例，$L$ 是 $d$ 的最小值。

3. Select

sknet-4.png

上圖右下角的式子看起來複雜但其實意思是：

  • 各分支在這邊皆有一個全連接層對 $z$ 做運算。所以兩個分支就有兩個全連接層，以上圖來說，$a$、$b$ 便是這兩個全連接層的產物。
  • 之後各分支的全連接層進行 channel-wise 的 softmax，讓各全連接層各對應的通道的和為 1，做到比例分配的效果。從上圖來看，可以看到a、b相對應的通道有一個是黑色另一個不是黑色，表示在這個通道上對不同 kernel size 的分支產生的結果的比例不同，也就是：縱向的（每一列）進行softmax[3]。
  • 之後依各分支的全連接層對各分支一開始的卷積結果進行權重計算最後相加得到結果 $V$。以上圖來說，也就是 $a$、$b$ 上各通道的值是原本卷積結果（$\widetilde U$、$\hat U$）各通道的權重。

與其它模型的比較

sknet-5.png

圖中 SK[M=......] 表示 SK 的超參數。

  • $M$：是分支數，也就是有幾種 kernel size。
  • $G$：是各分支的卷積層做分組卷積的分組數。
  • $r$： $\bold z$ 的維度為 $d=max(C/r,L)$，$r$ 是控制用的比例（$L$ 是 $d$ 的最小值）。

學習調節感受野並結合 SENet

感受野（receptive field），受到影響的區域。以下圖為例，做 3x3 的卷積時，Layer 3 的黃色格子受到 Layer 2 中間那 3x3 格影響；而這 3x3 格又是受到前一層 5x5 格的影響。用這個概念來思考，如果 kernel size 變大，感受野也會變大。SKNet 作者之一的李翔在他一篇文章[5]提到「SKNet 啟發自皮質神經元根據不同的刺激可動態調節其自身的感受野」。

sknet-6.png

圖片來源：[4]

SKNet 結合了 SENet 的想法，除了學習感受野的調節，也學習到不同通道間的關係。

範例程式碼

以下範例程式碼修改自 developer0hye 在 GitHub 上的 repository "SKNet-PyTorch" [6]。

class SKConv(nn.Module):
    def __init__(self, features, M=2, G=32, r=16, stride=1 ,L=32):
        """ Constructor
        Args:
            features: input channel dimensionality.
            M: the number of branchs.
            G: num of convolution groups.
            r: the ratio for compute d, the length of z.
            stride: stride, default 1.
            L: the minimum dim of the vector z in paper, default 32.
        """
        super(SKConv, self).__init__()
        d = max(int(features/r), L) # z 的維度
        self.M = M
        self.features = features
        self.convs = nn.ModuleList([]) # 各分支的卷積層
        for i in range(M):
            self.convs.append(nn.Sequential(
                nn.Conv2d(features, features, kernel_size=3, stride=stride, padding=1+i, dilation=1+i, groups=G, bias=False),
                nn.BatchNorm2d(features),
                nn.ReLU(inplace=False)
            ))
        self.gap = nn.AdaptiveAvgPool2d((1,1)) # Global average pooling
        self.fc = nn.Sequential(nn.Conv2d(features, d, kernel_size=1, stride=1, bias=False),
                                nn.BatchNorm2d(d),
                                nn.ReLU(inplace=False)) # Fuse 的全連結層
        self.fcs = nn.ModuleList([]) # 各分支加起來前的全連結層
        for i in range(M):
            self.fcs.append(
                 nn.Conv2d(d, features, kernel_size=1, stride=1)
            )
        self.softmax = nn.Softmax(dim=1) # Softmax

    def forward(self, x):
        batch_size = x.shape[0]
        # 各分支的卷積計算
        feats = [conv(x) for conv in self.convs]
        # U
        feats = torch.cat(feats, dim=1)
        feats = feats.view(batch_size, self.M, self.features, feats.shape[2], feats.shape[3])
        feats_U = torch.sum(feats, dim=1)
        # s
        feats_S = self.gap(feats_U)
        # z
        feats_Z = self.fc(feats_S)
        # 各分支的全連結層
        attention_vectors = [fc(feats_Z) for fc in self.fcs]
        # Softmax
        attention_vectors = torch.cat(attention_vectors, dim=1)
        attention_vectors = attention_vectors.view(batch_size, self.M, self.features, 1, 1)
        attention_vectors = self.softmax(attention_vectors)
        # 最後加起來
        feats_V = torch.sum(feats*attention_vectors, dim=1)
        return feats_V

Reference & Source

[1] X. Li, W. Wang, X. Hu and J. Yang, "Selective Kernel Networks," 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2019, pp. 510-519, doi: 10.1109/CVPR.2019.00060.

[2] https://github.com/vdumoulin/conv_arithmetic

[3] 引用自 CSDN 博主「ITOMG」的原創文章，遵循CC 4.0 BY-SA版權協議，轉載請附上原文出處連結及本聲明。原文連結：https://blog.csdn.net/ITOMG/article/details/89673593

[4] Lin, H.; Shi, Z.; Zou, Z. Maritime Semantic Labeling of Optical Remote Sensing Images with Multi-Scale Fully Convolutional Network. Remote Sens. 2017, 9, 480. https://doi.org/10.3390/rs9050480

[5] 李翔：SKNet——SENet孪生兄弟篇(https://zhuanlan.zhihu.com/p/59690223)

[6] Repository "SKNet-PyTorch" on GitHub by pppLang (https://github.com/developer0hye/SKNet-PyTorch)