在2017年末,Face++發了一篇論文[ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices ](https://arxiv.org/abs/1707.01083)討論了一個極有效率且可... ...
在2017年末,Face++發了一篇論文ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices討論了一個極有效率且可以運行在手機等移動設備上的網路結構——ShuffleNet。這個英文名我更願意翻譯成“重組通道網路”,ShuffleNet通過分組捲積與\(1 \times 1\)的捲積核來降低計算量,通過重組通道來豐富各個通道的信息。這個論文的mxnet源碼的開源地址為:MXShuffleNet。
分組捲積與核大小對計算量的影響
論文說中到“We propose using pointwise group convolutions to reduce computation complexity of 1 × 1 convolutions”,那麼為什麼用分組捲積與小的捲積核會減少計算的複雜度呢?先來看看捲積在編程中是如何實現的,Caffe與mxnet的CPU版本都是用差不多的方法實現的,但Caffe的計算代碼會更加簡潔。
不分組且只有一個樣本
在不分組與輸入的樣本量為1(batch_size=1)的條件下,輸出一個通道上的一個點是捲積核會與所有的通道捲積之積,如圖1所示:
在Caffe的計算方法中,先要將輸入張量為\(n \times C_{in} \times H_{in} \times W_{in}\)(n是batch_size)轉化為一個$ \left(C_{in} \times H_k \times H_w\right) \times \left(H_{in} \times W_{in}\right)\(的矩陣,這個過程叫**im2col**。最後得到的輸出張量為\)n \times C_{out} \times H_{in} \times W_{in}$。
得到的兩個矩陣Feature與Filter相乘得到輸出矩陣Output,再Reshape成\(C_{out} \times C_{in} \times H_k \times W_k\)張量:
\[
Filter_{C_{out} \times \left( C_{in} \times K_h \times K_w \right)} \times Feature_{\left(C_{in} \times H_k \times H_w\right) \times \left(H_{out} \times W_{out}\right)} = Output_{C_{out} \times (H_{out} \times W_{out})} \tag{1.1}
\]
現在的計算技術中,對方長度為\(n\)的方陣,計算量能從\(n^3\)代碼到\(n^{2.376}\),最小的複雜度現在仍然未知,本文為了方便計算量就以\(n^3\)為基準。所以式(1.1)的矩陣計算最普通的計算量\(Computation\)是:
\[
Computation=C_{out} \times H_{out} \times W_{out} \times \left( C_{in} \times K_h \times K_w \right)^2 \tag{1.2}
\]
從式(1.2)中可以看出來,捲積核的大小對計算量影響是很大的,\(3 \times 3\)的捲積核比\(1 \times 1\)的計算量要大\(3^4=81\)倍。
分組且只有一個樣本
什麼叫做分組,就是將輸入與輸出的通道分成幾組,比如輸出與輸入的通道數都是4個且分成2組,那第1、2通道的輸出只使用第1、2通道的輸入,同樣那第3、4通道的輸出只使用第1、2通道的輸入。也就是說,不同組的輸出與輸入沒有關係了,減少聯繫必然會使計算量減小,但同時也會導致信息的丟失。
當分成g組後,一層參數量的大小由\(Filter_{C_{out} \times \left( C_{in} \times K_h \times K_w \right)}\)變成\(Filter_{C_{out} \times \left( C_{in} \times K_h \times K_w / g \right)}\)。Feature Matrix的大小雖然沒發生變化,但是每一組的使用量是原來的$1/g,Filter也只用到所有參數的\(1/g\)\(。然後再迴圈計算\)g$次(同時FeatureMatrix與FilterMatrix要有地址偏移),那麼計算公式與計算量的大小為:
\[
Filter_{C_{out}/g \times \left( C_{in} \times K_h \times K_w /g \right)} \times Feature_{\left(C_{in} \times H_k \times H_w /g\right) \times \left(H_{out} \times W_{out}\right)} = Output_{C_{out}/g \times (H_{out} \times W_{out})} \tag{1.3}
\]
\[
Computation=C_{out} \times H_{out} \times W_{out} \times \left( C_{in} \times K_h \times K_w /g \right)^2 \tag{1.4}
\]
所以,分成\(g\)組可以使參數量變成原來的\(1/g\),計算量是原來的\(1/g^2\)。
多個樣本輸入
為了節省記憶體,多個樣本輸入的時候,上述的所有過程都不會改變,而是每一個樣本都運行一次上述的過程。
以上只是最簡單、粗略的分析,實際上計算效率的提升並不會有上述這麼多,一方面因為im2col會消耗與矩陣運算差不多的時間,另一方面因為現代的blas庫優化了矩陣運算,複雜度並沒有上述分析的那麼多,還有計算過程for迴圈是比較耗時的指令,即使用openmp也不能優化捲積的計算過程。
交換通道(Shuffle Channels)
在上面我提到過,分組會導致信息的丟失,那麼有沒有辦法來解決這個問題呢?這個論文給出的方法就是交換通道,因為在同一組中不同的通道蘊含的信息可能是相同的,如果在不同的組之後交換一些通道,那麼就能交換信息,使得各個組的信息更豐富,能提取到的特征自然就更多,這樣是有利於得到更好的結果。
ShuffleUnit
ShuffleUnit的設計參考了ResNet,總有兩個基本單元,兩人個基本單元功能不一樣,將他們組合起來就可以得到ShuffleNet。這樣的設計可以在增加網路的深度(比mobilenet深約一倍)的同時,減少參數總量與計算量(本人運行Cifar10時,速度大約是molibenet的10倍)。
源碼解讀
def combine(residual, data, combine):
if combine == 'add':
return residual + data
elif combine == 'concat':
return mx.sym.concat(residual, data, dim=1)
return None
add是代表圖6中的單元b),concat是代表圖6中的單元c)。
def channel_shuffle(data, groups):
data = mx.sym.reshape(data, shape=(0, -4, groups, -1, -2))
data = mx.sym.swapaxes(data, 1, 2)
data = mx.sym.reshape(data, shape=(0, -3, -2))
return data
這個函數就是交換通道的函數,函數的第一行data = mx.sym.reshape(data, shape=(0, -4, groups, -1, -2))是將輸入為\(n \times C_{in} \times H_{in} \times W_{in}\)reshape成\(n \times (C_{in}/g) \times g\times H_{in} \times W_{in}\),要註意的是mxnet中reshape不會改變張量在記憶體中的排列順序。至於要mxnet中的0,-1,-2,-3,-4的具體意義可以這樣看到:
import mxnet as mx
print(help(mx.sym.reshape))
可以看到輸出以下(只提取出一小部分,其餘的可用上述方法查看),這裡有各個參數的具體意義:
- ``0`` copy this dimension from the input to the output shape.
- ``-1`` infers the dimension of the output shape by using the remainder of the input dimensions
- ``-2`` copy all/remainder of the input dimensions to the output shape.
- ``-3`` use the product of two consecutive dimensions of the input shape as the output dimension.
- ``-4`` split one dimension of the input into two dimensions passed subsequent to -4 in shape (can contain -1).
函數的第二行是交換第一與第二個維度,那麼現在這個symbol的符號的shape就變成了\(n \times g \times (C_{in}/g) \times H_{in} \times W_{in}\)。這裡的第零個維度是\(n\)。要註意的是交換維度改變了張量在記憶體中的排列順序,改變了記憶體中的順序實現上就是完成了圖5c)中的Channel Shuffle操作,不同的顏色代碼數據在原來記憶體中的位置。
函數的最後一行合併了第一與第二個維度,輸出的張量與輸入的張量shape都是\(n \times C_{in} \times H_{in} \times W_{in}\)。
def shuffleUnit(residual, in_channels, out_channels, combine_type, groups=3, grouped_conv=True):
if combine_type == 'add':
DWConv_stride = 1
elif combine_type == 'concat':
DWConv_stride = 2
out_channels -= in_channels
first_groups = groups if grouped_conv else 1
bottleneck_channels = out_channels // 4
data = mx.sym.Convolution(data=residual, num_filter=bottleneck_channels,
kernel=(1, 1), stride=(1, 1), num_group=first_groups)
data = mx.sym.BatchNorm(data=data)
data = mx.sym.Activation(data=data, act_type='relu')
data = channel_shuffle(data, groups)
data = mx.sym.Convolution(data=data, num_filter=bottleneck_channels, kernel=(3, 3),
pad=(1, 1), stride=(DWConv_stride, DWConv_stride), num_group=groups)
data = mx.sym.BatchNorm(data=data)
data = mx.sym.Convolution(data=data, num_filter=out_channels,
kernel=(1, 1), stride=(1, 1), num_group=groups)
data = mx.sym.BatchNorm(data=data)
if combine_type == 'concat':
residual = mx.sym.Pooling(data=residual, kernel=(3, 3), pool_type='avg',
stride=(2, 2), pad=(1, 1))
data = combine(residual, data, combine_type)
return data
ShuffleUnit這個函數實現上是實現圖6的b)與c),add對應成b),comcat對應於c)。
def make_stage(data, stage, groups=3):
stage_repeats = [3, 7, 3]
grouped_conv = stage > 2
if groups == 1:
out_channels = [-1, 24, 144, 288, 567]
elif groups == 2:
out_channels = [-1, 24, 200, 400, 800]
elif groups == 3:
out_channels = [-1, 24, 240, 480, 960]
elif groups == 4:
out_channels = [-1, 24, 272, 544, 1088]
elif groups == 8:
out_channels = [-1, 24, 384, 768, 1536]
data = shuffleUnit(data, out_channels[stage - 1], out_channels[stage],
'concat', groups, grouped_conv)
for i in range(stage_repeats[stage - 2]):
data = shuffleUnit(data, out_channels[stage], out_channels[stage],
'add', groups, True)
return data
def get_shufflenet(num_classes=10):
data = mx.sym.var('data')
data = mx.sym.Convolution(data=data, num_filter=24,
kernel=(3, 3), stride=(2, 2), pad=(1, 1))
data = mx.sym.Pooling(data=data, kernel=(3, 3), pool_type='max',
stride=(2, 2), pad=(1, 1))
data = make_stage(data, 2)
data = make_stage(data, 3)
data = make_stage(data, 4)
data = mx.sym.Pooling(data=data, kernel=(1, 1), global_pool=True, pool_type='avg')
data = mx.sym.flatten(data=data)
data = mx.sym.FullyConnected(data=data, num_hidden=num_classes)
out = mx.sym.SoftmaxOutput(data=data, name='softmax')
return out
這兩個函數可以直接得到作者在論文中的表:
結果比較
論文後面用了種實驗證明這兩個技術的有效性,且證實了ShuffleNet的優秀,這裡就不細說,看論文後面的表就能一目瞭然。