MobileNet(v1/2)、ShuffleNet等の高速なモデルの構成要素と何故高速なのかの解説

はじめに

MobileNet系の高速なモデルアーキテクチャに利用される構成要素と、それらを利用したモデルについて、何故高速なのか観点と、空間方向の畳み込みとチャネル方向の畳み込みがどのようになされているかという観点で整理を行う。

高速なモデルアーキテクチャに利用される構成要素

まず、高速なモデルに利用される畳み込み構成要素について、計算量と、空間方向の畳み込みとチャネル方向の畳み込みがどのようになっているかを説明する。

まず、一般的な畳み込みの計算量を確認する。
入力特徴マップのサイズを$H \times W$、入力チャネル数を$N$、カーネルサイズを$K \times K$、出力チャネル数を$M$とすると、一般的な畳み込み層の計算量は、$H W N K^2 M$となる。
これは、入力特徴マップの1箇所につき畳み込みのコストが$K^2 N$で、これを入力特徴マップの$H W$箇所に適用することで、1チャネルの出力特徴マップが生成されることから、出力特徴マップが$M$個のケースでは$H W N K^2 M$の計算量になる。

自明ではあるが、ポイントは、計算量が、(1) 入力特徴マップの空間的なサイズ$H W$に比例する、(2) 畳み込みカーネルのサイズ$K^2$に比例する、(3) 入力特徴マップと出力特徴マップのチャネル数の積$N M$に比例する、ということである。

この計算量は、入力特徴マップの空間方向とチャネル方向どちらにも同時に畳み込みを行った場合であり、これをサボることで計算量を削減することが高速なモデルアーキテクチャの基本的な考え方である。

Convolution

まずは通常の畳み込みから。前述の通り、計算量は$H W N K^2 M$となる。

出力特徴マップの各値が、入力特徴マップの空間方向およびチャネル方向のどの範囲を畳み込まれた結果として得られるかという依存関係を可視化してみると、入力と出力の間を結んでいる線（重み）の数が計算量となる。

例えば、最も多用される畳み込みであるconv3x3は、上記のように空間方向、チャネル方向どちらにも畳み込みが行われる。空間方向は1次元のみを可視化しており、空間的に近い入力についてのみ畳み込みが行われることが分かる。

Pointwise convolutionとも呼ばれる、チャネル方向の次元を操作する（主に次元削減）際に用いられるconv1x1¹の場合は上図となる。
この畳み込みは、カーネルサイズが1x1であるため計算量が$H W N M$となり、例えばconv3x3と比較すると1/9の計算量となる。MobileNet系のアーキテクチャでは、チャネル方向の特徴を混ぜ合わせるために利用される。

Grouped Convolution

Grouped convolutionは、入力特徴マップをチャネル方向に分割し、その分割された各チャネルグループ内で畳み込みが行われる。定義上は、入出力のチャネル数が異なっていても良いが、同じチャネル数で利用されることが多い。
グループ数をGとすると、計算量は$H W N K^2 M / G$となり、グループ数が多いほど計算量が小さくなることが分かる。

グループ数が2のgrouped conv3x3。チャネル方向の畳み込みがグループ化されているので、通常の畳み込みと比較して計算量が小さいことが分かる。

グループ数が3のgrouped conv3x3。グループ数が2の場合と比較して更に計算量が小さくなる。

グループ数が2のgrouped conv1x1。カーネルサイズを1x1にすることもできる。もはや何をしているのかという感じだが、ShuffleNetで利用される。

グループ数が3のgrouped conv1x1。

Depthwise Convolution

Depthwise convolution^{2 3 4}は、入力特徴マップのチャネル毎に独立に畳み込みを行う。Grouped convolutionにおいて、入力チャネル数と出力チャネル数が同一で、グループ数が入力チャネル数と同じである特殊ケースとしても定義できる。
計算量は$H W N K^2$となり、通常の畳み込みの$1/M$ ($M=N$) の計算量となる。

上記のように、depthwise convolutionはチャネル方向には畳み込みを行わないことで、大幅な計算量削減を実現している。

Channel Shuffle

Channel ShuffleはShuffleNet⁵で提案された、入力特徴マップのチャネルの順番を入れ替える操作である。テンソルのreshapeとtransposeにより実現される。
具体的には、grouped convolutionと同時に利用する前提があり、そのグループ数を$G$、入力チャネル数を$G N' (=N)$とすると、チャネルの次元を ($G$, $N'$) にreshapeし、更にその2次元をtransposeし、flattenするという操作として定義される。
計算量は定義できないが、オーバーヘッドは存在するはずなので、計算量のFLOPsだけでの比較は危険である。

グループ数が3の場合のchannel shuffle。畳み込みは行われず、入力チャネルの順序の入れ替えだけが行われる。

グループ数が2の場合のchannel shuffle。

軽量なモデルアーキテクチャ

以降では、具体的なモデルアーキテクチャについて、何故高速なのか観点と、空間方向の畳み込みとチャネル方向の畳み込みがどのようになされているかを見ていく。

ResNet (Bottleneck構造)

高速なモデルアーキテクチャではないが、bottleneck構造のResNet⁶（のresidual unit）をベースとして、後のモデルアーキテクチャを比較すると分かりやすい。

Bottleneck構造のResNet（のresidual unit）は、上記のように、conv1x1、conv3x3、conv1x1を重ねたものである。最初のconv1x1により次元削減を行うことで、計算量の大きなconv3x3の計算量を低減している。

ResNeXt

ResNeXt⁷は、端的に言えば、上記のresidual unitにおいて、conv3x3をgrouped conv3x3にすることで高速化したモデルアーキテクチャである。実際には、最初のconv1x1の次元削減率を小さくし（ResNetと比較してあまり次元削減しない）、ResNetと同等の計算量となるように調整し、精度向上を実現している。

MobileNet (Separable Conv)

MobileNet⁸は、下記のように、depthwise convとconv1x1 (pointwise conv) を連続して行うモジュール (separable convolution) を積み重ねて構築されるモデルである。

通常の畳み込みが、空間方向とチャネル方向を同時に畳み込むのに対し、separable convはそれらを個別に行う（factorizeする）ことで、計算量を$H W N K^2 M$から$H W N K^2~(\mathrm{depthwise}) + H W N M~(\mathrm{conv1x1})$に削減している。比率で言うと、$1/K^2 + 1/M$であり、通常$M >> K^2$である (e.g. $K=3$, $M \ge 32$) ことから、計算量は1/8〜1/9程度に削減されることが分かる。
重要なポイントとして、計算量の観点では、conv1x1のほうがボトルネックになっていることが挙げられる。

ShuffleNet

ShuffleNetのスタート地点は、MobileNetのボトルネックがconv1x1であるという点である。そもそもconv1x1自体がカーネルサイズが1x1の軽量な畳み込みであるため、もはや計算量の削減は見込めないように見えるが、前述のgrouped conv1x1を利用して計算量の削減を実現している。

上記がShuffleNetを構成するモジュールである。
ここでのポイントは、grouped conv1x1を行った後のchannel shuffleである。Channel shuffleがない場合、grouped convのグループ化されたチャネル間での畳み込みが行われず、特徴抽出が非効率になり、精度が低下する。

グループ数を変えた場合のShuffleNetを構成するモジュール。

MobileNet-v2

MobileNet-v2⁹は、MobileNetのseparable convを、ResNetのbottleneck構造のように変更したモデルアーキテクチャである。

上記から分かるように、通常のbottleneck構造とは逆に、次元を増加させた後にdepthwise convを行い、その後次元を削減する形を取っている。

これは、上記のように、順序を入れ替えてMobileNet-v1 (separable conv) と比較してみると何をやっているかが分かりやすい。すなわち、（論文には明記されていないように見えるが）アイディアとしてはShuffleNetと同じで、conv1x1の計算量がボトルネックとなっているので、1つの計算量の大きなconv1x1を、計算量の小さなconv1x1を2つ利用することで近似しているということである。
具体的には、入出力チャネル数が$N$のconv1x1の計算量は$H W N^2$である。これに対し、入力チャネル数が$N$、出力チャネル数が$N/t$のconv1x1と、入力チャネル数が$N/t$、出力チャネル数が$N$のconv1x1の計算量の合計は$2 H W N^2 / t$となる。ここで、$t$はチャネルの拡張率 (expansion factor) である。
例えば文献⁵で利用されている$t=6$の場合、conv1x1の計算量が1/3になることが分かる。

FD-MobileNet

最後に、Fast-Downsampling MobileNet (FD-MobileNet)¹⁰を紹介する。
このモデルがどういうものかというと、MobileNetにおいて、画像サイズを1/2に縮小するダウンサンプリングの処理を、なるべく浅い層で行う（だけ）というものである。
実はこれだけで計算量を削減することができるのだが、その理由は近年のモデルアーキテクチャの設計方針と、通常の畳み込みとseparable convの計算量の違いにある。
すなわち、歴史的には、VGGNet以降のモデルにおいては、畳み込みを複数回行った後、特徴マップのサイズを1/2にダウンサンプリングしつつチャネル数を2倍にするという設計方針が一般的に採用されている。
通常の畳み込みの場合、計算量は$H W N K^2 M$であるため、ダウンサンプリングの前後で計算量が変化しない。
一方、separable convの計算量は$H W N (K^2 + M)$であり、チャネル数（$N$, $M$）が大きくなると相対的にカーネルサイズ影響される計算量$K^2$が無視できるようになるという特性がある。このため、fast-downsamplingを行うと、モデル全体の計算量が削減できるのである。

以上、MobileNet系のモデルアーキテクチャと、高速化のアイディアについて解説した。何故高速化ができるのかについては、論文にも明確な記載がないケースがあったので参考になれば幸いである。

References

1. M. Lin, Q. Chen, and S. Yan, "Network in Network," in Proc. of ICLR, 2014. ↩

2. L. Sifre, "Rigid-motion Scattering for Image Classification, Ph.D. thesis, 2014. ↩

3. L. Sifre and S. Mallat, "Rotation, Scaling and Deformation Invariant Scattering for Texture Discrimination," in Proc. of CVPR, 2013. ↩

4. F. Chollet, "Xception: Deep Learning with Depthwise Separable Convolutions," in Proc. of CVPR, 2017. ↩

5. X. Zhang, X. Zhou, M. Lin, and J. Sun, "ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices," in arXiv:1707.01083, 2017. ↩ ↩²

6. K. He, X. Zhang, S. Ren, and J. Sun, "Deep Residual Learning for Image Recognition," in Proc. of CVPR, 2016. ↩

7. S. Xie, R. Girshick, P. Dollár, Z. Tu, and K. He, "Aggregated Residual Transformations for Deep Neural Networks," in Proc. of CVPR, 2017. ↩

8. A. G. Howard, M. Zhu, B. Chen, D. Kalenichenko, W. Wang, T. Weyand, M. Andreetto, and H. Adam, "Mobilenets: Efficient Convolutional Neural Networks for Mobile Vision Applications," in arXiv:1704.04861, 2017. ↩

9. M. Sandler, A. Howard, M. Zhu, A. Zhmoginov, and L. Chen, "MobileNetV2: Inverted Residuals and Linear Bottlenecks," in arXiv:1801.04381v3, 2018. ↩

10. Z. Qin, Z. Zhang, X. Chen, and Y. Peng, "FD-MobileNet: Improved MobileNet with a Fast Downsampling Strategy," in arXiv:1802.03750, 2018. ↩