CNN模型之ShuffleNet

歡迎交流與轉載，文章會同步發布在公眾號：機器學習演算法全棧工程師(Jeemy110)

引言

ShuffleNet是曠視科技最近提出的一種計算高效的CNN模型，其和MobileNet和SqueezeNet等一樣主要是想應用在移動端。所以，ShuffleNet的設計目標也是如何利用有限的計算資源來達到最好的模型精度，這需要很好地在速度和精度之間做平衡。ShuffleNet的核心是採用了兩種操作：pointwise group convolution和channle shuffle，這在保持精度的同時大大降低了模型的計算量。目前移動端CNN模型主要設計思路主要是兩個方面：模型結構設計和模型壓縮。ShuffleNet和MobileNet一樣屬於前者，都是通過設計更高效的網路結構來實現模型變小和變快，而不是對一個訓練好的大模型做壓縮或者遷移。下面我們將詳細講述ShuffleNet的設計思路，網路結構及模型效果，最後使用Pytorch來實現ShuffleNet網路。

設計理念

ShuffleNet的核心設計理念是對不同的channels進行shuffle來解決group convolution帶來的弊端。Group convolution是將輸入層的不同特徵圖進行分組，然後採用不同的卷積核再對各個組進行卷積，這樣會降低卷積的計算量。因為一般的卷積都是在所有的輸入特徵圖上做卷積，可以說是全通道卷積，這是一種通道密集連接方式（channel dense connection）。而group convolution相比則是一種通道稀疏連接方式（channel sparse connection）。使用group convolution的網路如Xception，MobileNet，ResNeXt等。Xception和MobileNet採用了depthwise convolution，這其實是一種比較特殊的group convolution，因此此時分組數恰好等於通道數，意味著每個組只有一個特徵圖。但是這些網路存在一個很大的弊端是採用了密集的1x1卷積，或者說是dense pointwise convolution，這裡說的密集指的是卷積是在所有通道上進行的。所以，實際上比如ResNeXt模型中1x1卷積基本上佔據了93.4%的乘加運算。那麼不如也對1x1卷積採用channel sparse connection，那樣計算量就可以降下來了。但是group convolution存在另外一個弊端，如圖1-a所示，其中GConv是group convolution，這裡分組數是3。可以看到當堆積GConv層後一個問題是不同組之間的特徵圖是不通信的，這就好像分了三個互不相干的路，大家各走各的，這目測會降低網路的特徵提取能力。這樣你也可以理解為什麼Xception，MobileNet等網路採用密集的1x1卷積，因為要保證group convolution之後不同組的特徵圖之間的信息交流。但是達到上面那個目的，我們不一定非要採用dense pointwise convolution。如圖1-b所示，你可以對group convolution之後的特徵圖進行「重組」，這樣可以保證接下了採用的group convolution其輸入來自不同的組，因此信息可以在不同組之間流轉。這個操作等價於圖2-c，即group convolution之後對channles進行shuffle，但並不是隨機的，其實是「均勻地打亂」。在程序上實現channle shuffle是非常容易的：假定將輸入層分為 組，總通道數為 ，首先你將通道那個維度拆分為 兩個維度，然後將這兩個維度轉置變成 ，最後重新reshape成一個維度。如果你不太理解這個操作，你可以試著動手去試一下，發現僅需要簡單的維度操作和轉置就可以實現均勻的shuffle。利用channle shuffle就可以充分發揮group convolution的優點，而避免其缺點。

網路結構

基於上面的設計理念，首先來構造ShuffleNet的基本單元，如圖2所示。ShuffleNet的基本單元是在一個殘差單元的基礎上改進而成的。如圖2-a所示，這是一個包含3層的殘差單元：首先是1x1卷積，然後是3x3的depthwise convolution（DWConv，主要是為了降低計算量），這裡的3x3卷積是瓶頸層（bottleneck），緊接著是1x1卷積，最後是一個短路連接，將輸入直接加到輸出上。現在，進行如下的改進：將密集的1x1卷積替換成1x1的group convolution，不過在第一個1x1卷積之後增加了一個channle shuffle操作。值得注意的是3x3卷積後面沒有增加channle shuffle，按paper的意思，對於這樣一個殘差單元，一個channle shuffle操作是足夠了。還有就是3x3的depthwise convolution之後沒有使用ReLU激活函數。改進之後如圖2-b所示。對於殘差單元，如果stride=1時，此時輸入與輸出shape一致可以直接相加，而當stride=2時，通道數增加，而特徵圖大小減小，此時輸入與輸出不匹配。一般情況下可以採用一個1x1卷積將輸入映射成和輸出一樣的shape。但是在ShuffleNet中，卻採用了不一樣的策略，如圖2-c所示：對原輸入採用stride=2的3x3 avg pool，這樣得到和輸出一樣大小的特徵圖，然後將得到特徵圖與輸出進行連接（concat），而不是相加。這樣做的目的主要是降低計算量與參數大小。

基於上面改進的ShuffleNet基本單元，設計的ShuffleNet模型如表1所示。可以看到開始使用的普通的3x3的卷積和max pool層。然後是三個階段，每個階段都是重複堆積了幾個ShuffleNet的基本單元。對於每個階段，第一個基本單元採用的是stride=2，這樣特徵圖width和height各降低一半，而通道數增加一倍。後面的基本單元都是stride=1，特徵圖和通道數都保持不變。對於基本單元來說，其中瓶頸層，就是3x3卷積層的通道數為輸出通道數的1/4，這和殘差單元的設計理念是一樣的。不過有個細節是，對於stride=2的基本單元，由於原輸入會貢獻一部分最終輸出的通道數，那麼在計算1/4時到底使用最終的通道數，還是僅僅未concat之前的通道數。文章沒有說清楚，但是個人認為應該是後者吧。其中$g$控制了group convolution中的分組數，分組越多，在相同計算資源下，可以使用更多的通道數，所以 越大時，採用了更多的卷積核。這裡給個例子，當 時，對於第一階段的第一個基本單元，其輸入通道數為24，輸出通道數為240，但是其stride=2，那麼由於原輸入通過avg pool可以貢獻24個通道，所以相當於左支只需要產生240-24=216通道，中間瓶頸層的通道數就為216/4=54。其他的可以以此類推。當完成三階段後，採用global pool將特徵圖大小降為1x1，最後是輸出類別預測值的全連接層。

模型效果

那麼ShuffleNet的模型效果如何呢？表2給出了採用不同的$g$值的ShuffleNet在ImageNet上的實驗結果。可以看到基本上當$g$越大時，效果越好，這是因為採用更多的分組後，在相同的計算約束下可以使用更多的通道數，或者說特徵圖數量增加，網路的特徵提取能力增強，網路性能得到提升。注意Shuffle 1x是基準模型，而0.5x和0.25x表示的是在基準模型上將通道數縮小為原來的0.5和0.25。

除此之外，作者還對比了不採用channle shuffle和採用之後的網路性能對比，如表3所示。可以清楚的看到，採用channle shuffle之後，網路性能更好，從而證明channle shuffle的有效性。

然後是ShuffleNet與MobileNet的對比，如表4所示。可以看到ShuffleNet不僅計算複雜度更低，而且精度更好。

ShuffleNet與其他CNN網路的對比可以去原始paper中更深入的了解。

ShuffleNet的Pytorch實現

這裡我們使用Pytorch來實現ShuffleNet，Pytorch是Facebook提出的一種深度學習動態框架，之所以採用Pytorch是因為其nn.Conv2d天生支持group convolution，不過儘管TensorFlow不支持直接的group convolution，但是其實可以自己間接地來實現。不過患有懶癌的我還是使用Pytorch吧。

首先我們來實現channle shuffle操作，就按照前面講述的思路來實現：

def shuffle_channels(x, groups):n """shuffle channels of a 4-D Tensor"""n batch_size, channels, height, width = x.size()n assert channels % groups == 0n channels_per_group = channels // groupsn # split into groupsn x = x.view(batch_size, groups, channels_per_group,n height, width)n # transpose 1, 2 axisn x = x.transpose(1, 2).contiguous()n # reshape into orignaln x = x.view(batch_size, channels, height, width)n return xn

然後我們實現ShuffleNet中stride=1的基本單元：

class ShuffleNetUnitA(nn.Module):n """ShuffleNet unit for stride=1"""n def __init__(self, in_channels, out_channels, groups=3):n super(ShuffleNetUnitA, self).__init__()n assert in_channels == out_channelsn assert out_channels % 4 == 0n bottleneck_channels = out_channels // 4n self.groups = groupsn self.group_conv1 = nn.Conv2d(in_channels, bottleneck_channels,n 1, groups=groups, stride=1)n self.bn2 = nn.BatchNorm2d(bottleneck_channels)n self.depthwise_conv3 = nn.Conv2d(bottleneck_channels,n bottleneck_channels,n 3, padding=1, stride=1,n groups=bottleneck_channels)n self.bn4 = nn.BatchNorm2d(bottleneck_channels)n self.group_conv5 = nn.Conv2d(bottleneck_channels, out_channels,n 1, stride=1, groups=groups)n self.bn6 = nn.BatchNorm2d(out_channels)nn def forward(self, x):n out = self.group_conv1(x)n out = F.relu(self.bn2(out))n out = shuffle_channels(out, groups=self.groups)n out = self.depthwise_conv3(out)n out = self.bn4(out)n out = self.group_conv5(out)n out = self.bn6(out)n out = F.relu(x + out)n return outn

然後是中stride=2的基本單元：

class ShuffleNetUnitB(nn.Module):n """ShuffleNet unit for stride=2"""n def __init__(self, in_channels, out_channels, groups=3):n super(ShuffleNetUnitB, self).__init__()n out_channels -= in_channelsn assert out_channels % 4 == 0n bottleneck_channels = out_channels // 4n self.groups = groupsn self.group_conv1 = nn.Conv2d(in_channels, bottleneck_channels,n 1, groups=groups, stride=1)n self.bn2 = nn.BatchNorm2d(bottleneck_channels)n self.depthwise_conv3 = nn.Conv2d(bottleneck_channels,n bottleneck_channels,n 3, padding=1, stride=2,n groups=bottleneck_channels)n self.bn4 = nn.BatchNorm2d(bottleneck_channels)n self.group_conv5 = nn.Conv2d(bottleneck_channels, out_channels,n 1, stride=1, groups=groups)n self.bn6 = nn.BatchNorm2d(out_channels)nn def forward(self, x):n out = self.group_conv1(x)n out = F.relu(self.bn2(out))n out = shuffle_channels(out, groups=self.groups)n out = self.depthwise_conv3(out)n out = self.bn4(out)n out = self.group_conv5(out)n out = self.bn6(out)n x = F.avg_pool2d(x, 3, stride=2, padding=1)n out = F.relu(torch.cat([x, out], dim=1))n return outn

最後是 的ShuffleNet的實現：

class ShuffleNet(nn.Module):n """ShuffleNet for groups=3"""n def __init__(self, groups=3, in_channels=3, num_classes=1000):n super(ShuffleNet, self).__init__()nn self.conv1 = nn.Conv2d(in_channels, 24, 3, stride=2, padding=1)n stage2_seq = [ShuffleNetUnitB(24, 240, groups=3)] + n [ShuffleNetUnitA(240, 240, groups=3) for i in range(3)]n self.stage2 = nn.Sequential(*stage2_seq)n stage3_seq = [ShuffleNetUnitB(240, 480, groups=3)] + n [ShuffleNetUnitA(480, 480, groups=3) for i in range(7)]n self.stage3 = nn.Sequential(*stage3_seq)n stage4_seq = [ShuffleNetUnitB(480, 960, groups=3)] + n [ShuffleNetUnitA(960, 960, groups=3) for i in range(3)]n self.stage4 = nn.Sequential(*stage4_seq)n self.fc = nn.Linear(960, num_classes)nn def forward(self, x):n net = self.conv1(x)n net = F.max_pool2d(net, 3, stride=2, padding=1)n net = self.stage2(net)n net = self.stage3(net)n net = self.stage4(net)n net = F.avg_pool2d(net, 7)n net = net.view(net.size(0), -1)n net = self.fc(net)n logits = F.softmax(net)n return logitsn

完整實現可以參見GitHub。

總結

本文主要介紹了ShuffleNet的核心設計思路以及網路架構，最後使用Pytorch來實現。說點題外話，在之前計算力不足時，CNN模型有時會採用group convolution，而隨著計算力的提升，目前大部分的CNN採用dense channle connections，但是現在一些研究又轉向了group convolution來提升速度，這有點戲劇化。不過ShuffleNet通過channle shuffle這一個trick解決了group convolution的一個副作用，還是值得稱讚的。

參考文獻

1. ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices

歡迎交流與轉載，文章會同步發布在公眾號：機器學習演算法全棧工程師(Jeemy110)