SCConv：即插即用的空間和通道重建卷積

本文作者提出了一種名為SCConv（Spatial and Channel reconstruction Convolution, 空間和通道重建卷積）的卷積模塊，目的是減少卷積神經(jīng)網(wǎng)絡(luò)中特征之間的空間和通道冗余，從而壓縮CNN模型并提高其性能。

作者設(shè)計(jì)的 SCConv 模塊，包含兩個(gè)單元。一個(gè)名為SRU (Spatial Reconstruction Unit, 空間重構(gòu)單元)，一個(gè)名為CRU (Channel Reconstruction Unit, 通道重構(gòu)單元)。其中 SRU 通過(guò)分離-重構(gòu)方法來(lái)減少空間冗余，CRU 則使用分割-轉(zhuǎn)換-融合方法來(lái)減少通道冗余。這兩個(gè)單元協(xié)同工作，以減少CNN中特征的冗余信息。

作者指出，SCConv 是一種可以直接替代標(biāo)準(zhǔn)卷積操作的插件式卷積模塊 ，可以應(yīng)用于各種卷積神經(jīng)網(wǎng)絡(luò)中，從而降低冗余特征并減少計(jì)算復(fù)雜性。

在后續(xù)的實(shí)驗(yàn)中，文章作者認(rèn)為相對(duì)于其他流行的 SOTA 方法，他們提出的 SCConv 可以以更低的計(jì)算成本獲得更高的準(zhǔn)確率。下圖是 ResNet50 在 ImageNet 上的 Top1 準(zhǔn)確性測(cè)試結(jié)果。

SCConv

如下圖，SCConv 由兩個(gè)單元組成，即空間重構(gòu)單元 (SRU) 和信道重構(gòu)單元 (CRU) ，兩個(gè)單元按順序排列。輸入的特征 X 先經(jīng)過(guò)空間重構(gòu)單元，得到空間細(xì)化的特征Xw。再經(jīng)過(guò)通道重構(gòu)單元，得到通道提煉的特征 Y 作為輸出。

SCConv 模塊利用了特征之間的空間冗余和信道冗余，模塊可以無(wú)縫集成到任何 CNN 框架中，減少特征之間的冗余，提高 CNN 特征的代表性。

作者對(duì)SRU和CRU進(jìn)行不同的組合，包括：

1. 不使用 SRU 和 CRU

2. 單獨(dú)使用 SRU

3. 單獨(dú)使用 CRU

4. 并行使用 SRU 和 CRU

5. 先使用 CRU 再使用 SRU

6. 先使用 SRU 在使用 CRU

最終發(fā)現(xiàn)先使用 SRU 再使用 CRU 的效果最好。

下面詳細(xì)介紹 SRU 和 CRU 這兩個(gè)單元。

SRU 空間重建單元

在作者的設(shè)計(jì)中，該單元采用分離-重構(gòu)的方法。

分離操作的目的是將信息量大的特征圖從信息量小的特征圖中分離出來(lái)，與空間內(nèi)容相對(duì)應(yīng)。作者使用組歸一化 (Group Normalization) 里的縮放因子來(lái)評(píng)估不同特征圖中的信息含量。

經(jīng)過(guò) SRU 處理后，信息量大的特征從信息量小的特征中分離出來(lái)，減少了空間維度上的冗余特征。

CRU 通道重建單元

在作者的設(shè)計(jì)中，該單元采用分割-轉(zhuǎn)換-融合的方法。

實(shí)驗(yàn)

消融實(shí)驗(yàn)

下圖的消融實(shí)驗(yàn)確定了 SRU 和 CRU 的排列方式

下圖的消融實(shí)驗(yàn)確定了 CRU 中的拆分系數(shù)α

圖片分類(lèi)實(shí)驗(yàn)

下圖是與其他 SOTA 方法的比較，作者認(rèn)為在所有的情況下，SCConv-embedded 模型的準(zhǔn)確性都優(yōu)于先前所有的網(wǎng)絡(luò)。在某些模型中，對(duì)比同類(lèi)模型在減少參數(shù)和 FLOPs 的同時(shí)還實(shí)現(xiàn)了更高的準(zhǔn)確率

相關(guān)代碼中文注釋

import torch  # 導(dǎo)入 PyTorch 庫(kù)
import torch.nn.functional as F  # 導(dǎo)入 PyTorch 的函數(shù)庫(kù)
import torch.nn as nn  # 導(dǎo)入 PyTorch 的神經(jīng)網(wǎng)絡(luò)模塊

# 自定義 GroupBatchnorm2d 類(lèi)，實(shí)現(xiàn)分組批量歸一化
class GroupBatchnorm2d(nn.Module):
    def __init__(self, c_num:int, group_num:int = 16, eps:float = 1e-10):
        super(GroupBatchnorm2d,self).__init__()  # 調(diào)用父類(lèi)構(gòu)造函數(shù)
        assert c_num >= group_num  # 斷言 c_num 大于等于 group_num
        self.group_num  = group_num  # 設(shè)置分組數(shù)量
        self.gamma      = nn.Parameter(torch.randn(c_num, 1, 1))  # 創(chuàng)建可訓(xùn)練參數(shù) gamma
        self.beta       = nn.Parameter(torch.zeros(c_num, 1, 1))  # 創(chuàng)建可訓(xùn)練參數(shù) beta
        self.eps        = eps  # 設(shè)置小的常數(shù) eps 用于穩(wěn)定計(jì)算

    def forward(self, x):
        N, C, H, W  = x.size()  # 獲取輸入張量的尺寸
        x           = x.view(N, self.group_num, -1)  # 將輸入張量重新排列為指定的形狀
        mean        = x.mean(dim=2, keepdim=True)  # 計(jì)算每個(gè)組的均值
        std         = x.std(dim=2, keepdim=True)  # 計(jì)算每個(gè)組的標(biāo)準(zhǔn)差
        x           = (x - mean) / (std + self.eps)  # 應(yīng)用批量歸一化
        x           = x.view(N, C, H, W)  # 恢復(fù)原始形狀
        return x * self.gamma + self.beta  # 返回歸一化后的張量

# 自定義 SRU（Spatial and Reconstruct Unit）類(lèi)
class SRU(nn.Module):
    def __init__(self,
                 oup_channels:int,  # 輸出通道數(shù)
                 group_num:int = 16,  # 分組數(shù)，默認(rèn)為16
                 gate_treshold:float = 0.5,  # 門(mén)控閾值，默認(rèn)為0.5
                 torch_gn:bool = False  # 是否使用PyTorch內(nèi)置的GroupNorm，默認(rèn)為False
                 ):
        super().__init__()  # 調(diào)用父類(lèi)構(gòu)造函數(shù)

         # 初始化 GroupNorm 層或自定義 GroupBatchnorm2d 層
        self.gn = nn.GroupNorm(num_channels=oup_channels, num_groups=group_num) if torch_gn else GroupBatchnorm2d(c_num=oup_channels, group_num=group_num)
        self.gate_treshold  = gate_treshold  # 設(shè)置門(mén)控閾值
        self.sigomid        = nn.Sigmoid()  # 創(chuàng)建 sigmoid 激活函數(shù)

    def forward(self, x):
        gn_x        = self.gn(x)  # 應(yīng)用分組批量歸一化
        w_gamma     = self.gn.gamma / sum(self.gn.gamma)  # 計(jì)算 gamma 權(quán)重
        reweights   = self.sigomid(gn_x * w_gamma)  # 計(jì)算重要性權(quán)重

        # 門(mén)控機(jī)制
        info_mask    = reweights >= self.gate_treshold  # 計(jì)算信息門(mén)控掩碼
        noninfo_mask = reweights < self.gate_treshold  # 計(jì)算非信息門(mén)控掩碼
        x_1          = info_mask * x  # 使用信息門(mén)控掩碼
        x_2          = noninfo_mask * x  # 使用非信息門(mén)控掩碼
        x            = self.reconstruct(x_1, x_2)  # 重構(gòu)特征
        return x

    def reconstruct(self, x_1, x_2):
        x_11, x_12 = torch.split(x_1, x_1.size(1) // 2, dim=1)  # 拆分特征為兩部分
        x_21, x_22 = torch.split(x_2, x_2.size(1) // 2, dim=1)  # 拆分特征為兩部分
        return torch.cat([x_11 + x_22, x_12 + x_21], dim=1)  # 重構(gòu)特征并連接

# 自定義 CRU（Channel Reduction Unit）類(lèi)
class CRU(nn.Module):
    def __init__(self, op_channel:int, alpha:float = 1/2, squeeze_radio:int = 2, group_size:int = 2, group_kernel_size:int = 3):
        super().__init__()  # 調(diào)用父類(lèi)構(gòu)造函數(shù)

        self.up_channel     = up_channel = int(alpha * op_channel)  # 計(jì)算上層通道數(shù)
        self.low_channel    = low_channel = op_channel - up_channel  # 計(jì)算下層通道數(shù)
        self.squeeze1       = nn.Conv2d(up_channel, up_channel // squeeze_radio, kernel_size=1, bias=False)  # 創(chuàng)建卷積層
        self.squeeze2       = nn.Conv2d(low_channel, low_channel // squeeze_radio, kernel_size=1, bias=False)  # 創(chuàng)建卷積層

        # 上層特征轉(zhuǎn)換
        self.GWC            = nn.Conv2d(up_channel // squeeze_radio, op_channel, kernel_size=group_kernel_size, stride=1, padding=group_kernel_size // 2, groups=group_size)  # 創(chuàng)建卷積層
        self.PWC1           = nn.Conv2d(up_channel // squeeze_radio, op_channel, kernel_size=1, bias=False)  # 創(chuàng)建卷積層

        # 下層特征轉(zhuǎn)換
        self.PWC2           = nn.Conv2d(low_channel // squeeze_radio, op_channel - low_channel // squeeze_radio, kernel_size=1, bias=False)  # 創(chuàng)建卷積層
        self.advavg         = nn.AdaptiveAvgPool2d(1)  # 創(chuàng)建自適應(yīng)平均池化層

    def forward(self, x):
        # 分割輸入特征
        up, low = torch.split(x, [self.up_channel, self.low_channel], dim=1)
        up, low = self.squeeze1(up), self.squeeze2(low)

        # 上層特征轉(zhuǎn)換
        Y1 = self.GWC(up) + self.PWC1(up)

        # 下層特征轉(zhuǎn)換
        Y2 = torch.cat([self.PWC2(low), low], dim=1)

        # 特征融合
        out = torch.cat([Y1, Y2], dim=1)
        out = F.softmax(self.advavg(out), dim=1) * out
        out1, out2 = torch.split(out, out.size(1) // 2, dim=1)
        return out1 + out2

# 自定義 ScConv（Squeeze and Channel Reduction Convolution）模型
class ScConv(nn.Module):
    def __init__(self, op_channel:int, group_num:int = 16, gate_treshold:float = 0.5, alpha:float = 1/2, squeeze_radio:int = 2, group_size:int = 2, group_kernel_size:int = 3):
        super().__init__()  # 調(diào)用父類(lèi)構(gòu)造函數(shù)

        self.SRU = SRU(op_channel, group_num=group_num, gate_treshold=gate_treshold)  # 創(chuàng)建 SRU 層
        self.CRU = CRU(op_channel, alpha=alpha, squeeze_radio=squeeze_radio, group_size=group_size, group_kernel_size=group_kernel_size)  # 創(chuàng)建 CRU 層

    def forward(self, x):
        x = self.SRU(x)  # 應(yīng)用 SRU 層
        x = self.CRU(x)  # 應(yīng)用 CRU 層
        return x

if __name__ == '__main__':
    x       = torch.randn(1, 32, 16, 16)  # 創(chuàng)建隨機(jī)輸入張量
    model   = ScConv(32)  # 創(chuàng)建 ScConv 模型
    print(model(x).shape)  # 打印模型輸出的形狀