CNNのボトルネック層（1x1畳み込み）による計算効率向上を理解する

「1x1畳み込みを使うと計算効率がよくなるよ」という話は聞いたことあっても、具体的にどれだけ良くなるのかを理論的に議論されることはあまり多くないように思います。自分の理解も含めて、一度これをちゃんと整理したいと思います。少し数式多めですがご勘弁ください。

実験のコードはこちらにあります。　https://github.com/koshian2/Inception-bottleneck

Summary

• 入力チャンネルがc, 出力がf, $c=\beta f$で、ボトルネックでのチャンネル圧縮が$\frac{f}{\alpha}$のとき、ボトルネックなしの場合と比べて、ボトルネック→Convの場合は$\alpha\beta$倍（BC比）、ボトルネック→Conv→ボトルネックの場合は$\alpha^2\beta$倍（BCB比）最大で計算回数を減らせる。
• しかし、実際は計算回数を1/10に減らせても、訓練速度は10倍にはらない。速くなるかどうかはモデルが深くなることに対するオーバーヘッドもあるので、フレームワーク次第。
• 極端に大きなαを使わなければ、ボトルネックを使うことで、モデルの高速化・軽量化とReLUの非線形性による精度の向上を両立できる。

Conv2Dの計算量

ごく普通のConv2Dレイヤーでどれだけの計算が行われているのかを考えます。

k×kのピクセルについて、1チャンネルあたりの計算回数は$k^2$なので、入力チャンネル数をかけた$k^2c$分が、1フィルター1回分の計算回数となります。これを横幅×縦幅（w×h）×フィルター数（f）回やるので、**Conv2Dの計算量は、$k^2c\cdot hwf$**となります。

以下話を簡単に進めるために、縦と横の解像度が等しいという仮定をおきます（大抵正方形の状態で畳み込みするので）。$h=w$という式により、

$$ k^2w^2cf \tag{1}$$

という計算回数の公式が成り立ちます。

1x1畳み込みの使い方

ボトルネック層（1x1畳み込み）を使った効率化のパターンは主に2パターンあるといえるでしょう。

1.ボトルネック→畳込み

まずは畳み込みの前にボトルネックを入れるパターン。これはInceptionなど目にすることが多いパターンです。

ここでボトルネック層のフィルター数を$\frac{f}{\alpha}$という、パラメーターαを用いて表すことにします。例えばf=32、ボトルネックのフィルター数が16なら、α＝2となります。

2.ボトルネック→畳込み→ボトルネック

こちらはあまり目にすることが少ないパターンなのですが、理論的には1と比べて強い計算効率化が出ます。例えばWide ResNetの論文の中では1つのパターンとして検討されていたものです。

計算コストが重いk×k畳み込みを、前後から1x1畳み込みで挟みます。前のボトルネックは次元削減、後ろのボトルネックは次元増加の効果があります。

つまりボトルネックなしも含めて3パターン存在することがわかります。

1.　ボトルネックなし（Conv）
2.　ボトルネック→畳込み（Bottleneck→Conv）
3.　ボトルネック→畳込み→ボトルネック（Bottleneck→Conv→Bottleneck）

この3パターンの計算回数を求めていきます。

パターンごとの計算回数の導出

1.ボトルネックなし

式(1)より、

$$k^2w^2cf\tag{2}$$

です。

2.ボトルネック→畳込み（Bottleneck→Conv）

ボトルネックは、カーネル1、入力チャンネル$c$、出力チャンネル$\frac{f}{\alpha}$なので、式(1)を適用すると、$w^2c\frac{f}{\alpha}$です。

本番の畳み込みは、カーネルk、入力チャンネル$\frac{f}{\alpha}$、出力チャンネル$f$なので、同様に計算すると$k^2w^2\frac{f}{\alpha}f$です。この2つを足すと、

$$w^2c\frac{f}{\alpha}+k^2w^2\frac{f}{\alpha}f=\frac{w^2f}{\alpha}(c+k^2f) \tag{3}$$

となります。ここで、パターン1（式2）をパターン2（式3）で割った値をBC比と呼ぶことにしましょう。例えば、パターン1のときに1億回計算していて、パターン2のときに1000万回の計算だったら、BC比は10となります。つまり、このBC比こそがボトルネックを組み込んだことによる計算効率の向上で、BC比が大きいほど計算効率が改善しているということがわかります。

式(2)を式(3)で割りましょう。

$$ BC=\frac{k^2w^2cf}{\frac{w^2f}{\alpha}(c+k^2f)}=\frac{\alpha k^2 c}{c+k^2 f}\tag{4}$$

3.ボトルネック→畳込み→ボトルネック（Bottleneck→Conv→Bottleneck）

今度はk×kの畳み込みが$\frac{f}{\alpha}$チャンネルであることに注意して同様に計算します。

$$w^2c\frac{f}{\alpha}+k^2w^2\frac{f}{\alpha}\frac{f}{\alpha}+w^2\frac{f}{\alpha}f = \frac{w^2f}{\alpha^2}(k^2f+\alpha(c+f)) \tag{5}$$

同様に、パターン1（式2）をパターン3（式5）で割った値をBCB比とします。BCB比もBC比と同じく、高いほど計算の効率化がされているということになります。

$$ BCB = \frac{k^2w^2cf}{\frac{w^2f}{\alpha^2}(k^2f+\alpha(c+f))}= \frac{\alpha^2k^2c}{\alpha(c+f)+k^2f}\tag{6}$$

BC比、BCB比の整理

ここで、出力チャンネル数fが入力チャンネル数cの定数倍で表されるとします。つまり、$c=\beta f$なるβが存在しているものとします。例えば、入力チャンネル数が128で出力チャンネル数が16なら、β＝8です。

BC比、BCB比のcをfで置き換えると、

$$BC = \frac{\alpha\beta k^2f}{\beta f+k^2 f}=\frac{\alpha\beta k^2}{\beta+k^2}\tag{7}$$
$$BCB = \frac{\alpha^2\beta k^2f}{\alpha\beta f+\alpha f + k^2f}=\frac{\alpha^2\beta k^2}{\alpha(\beta+1)+k^2} \tag{8}$$

このBC比とBCB比の具体的な値は後で確認しますが、k、つまりカーネルサイズが大きくなるほどある値に向かって収束するようになります。BC比とBCB比の極限を取りましょう。

\begin{align}
\lim_{k\to\infty}BC &= \alpha\beta \tag{9} \\
\lim_{k\to\infty}BCB &= \alpha^2\beta \tag{10}
\end{align}

かなりすっきりした式になりましたね。つまり、ボトルネック層による計算効率の向上は、極限的にはボトルネック層のフィルター数$\frac{f}{\alpha}$におけるαと、入力チャンネル数と出力チャンネル数を表すβに集約されるということになります。

BC比、BCB比のグラフ

BC比、BCB比をグラフで書いてみます。極限の値を点線で描きました。これで先程の極限の意味が理解しやすくなるかと思います。

横軸にk（カーネルサイズ）を取り、BC比とBCB比をαとβを変えてプロットしてみました。カーネルサイズが大きくなるほど式(9),(10)の極限に収束していくのがわかりますね。BC比はk=5を越えるとだいたい極限付近まで収束してしまうので、極端な話、極限の値を考えればよいということになります。

ここで、一度整理しましょう。BC比とは1x1畳み込みなし（ボトルネックなし）の場合の計算量を、ボトルネック→畳込み（Bottleneck→Conv）の計算量で割ったもので、ボトルネックを入れることによる計算量の削減度合いを表すバロメーターでした。そしてBC比はカーネルサイズkが十分に大きくなれば、$\alpha\beta$で表されます。同様にBCB比は、ボトルネックなしの場合を、ボトルネック→畳込み→ボトルネック（Bottleneck→Conv→Bottleneck）の場合で割った値で、BCB比はカーネルサイズkが十分に大きくなれば、$\alpha^2\beta$で表されます。あとはα、βの意味を解釈するだけです。

理論的なまとめ

ここまでの理論的な話をまとめます。

• BC比は$\alpha\beta$、BCB比は$\alpha^2\beta$に収束する。αはボトルネック層でのチャンネルの圧縮率なので、ボトルネック層で強く圧縮すればするほど計算効率は上がる。ただし128とかあまりに大きいαを使うとネットワークの表現力を損なうので、ほどほどに。
• βは入力チャンネルを出力チャンネルで割ったものである。つまり、出力チャンネルが入力チャンネルに比べて小さくなるケースでは、ボトルネックの効果が出る。Inceptionのようにいろいろなカーネルを、少ないチャンネル数で組み合わせる場合は特に効果が出やすい

つまり、Inceptionのあの構造はかなり考えられたものだったんですね。さすがはGoogleというところでしょうか。

1点目の「ほどほどに」ですが、α＝2とか常識的な値を使う限りでは、おそらくボトルネック層でのReLUの非線形性が、圧縮によるパフォーマンスの劣化を打ち消して、結果的に良い精度をもたらすと思われます。ここは後で実際に確認してみます。

Inceptionモジュール

これはGoogLeNetの論文に示されているInceptionモジュール3aを参考に一般化したものです（本当はのIncpetion3aは3x3の次元削減が1/2ではなく3/4です）。Inceptionモジュール自体が、ネットワークの配置される場所によって分岐のチャンネル分配がまちまちなので、議論しやすいようにこのように一般化しました。

この図での次元削減の「α＝2」です。また、βは$c=\beta f$で表します。

Inceptionモジュール＋AlexNet

実験ではこのInceptionモジュールをAlexNetの構成で実装しました。本物のInceptionだと時間かかりすぎるので訓練時間を軽減するためです。AlexNetの畳み込みをInceptionモジュールに置き換えたイメージです。

モジュール	解像度	入力チャンネル	出力チャンネル
Inception 1	32	3	96
Pooling 1	16	96	96
Inception 2	16	96	256
Pooling 2	8	256	256
Inception 3	8	256	384
Inception 4	8	384	384
Inception 5	8	384	256

このあとはGlobalAveragePoolingで即Softmaxに食わせます。

実験設定

このCNNを使ってCIFAR-10を分類するのですが、このInceptionモジュールの「αとモード」を変えて実験します。モードの設定は以下の3種類とします。3x3畳み込みと5x5畳み込みのブロックについて、

1. ボトルネック（1x1畳み込み）を使わない
2. ボトルネック→本来の畳み込みをする。ボトルネックでは「f/α」にチャンネルを圧縮し、続く畳込みでチャンネルをfにする。
3. ボトルネック→畳込み→ボトルネックとする。最初のボトルネックで「f/α」にチャンネルを圧縮し、続く畳み込みは「f/α」のまま、最後のボトルネックでチャンネルをfにする。

設定は理論編と同じです。以下の5パターン用意します。

パターン	モード	α	計算量	パラメーター数	計÷パ
A	1	-	231,091,200	2,960,874	78.0
B	2	2	87,118,848	839,838	103.7
C	2	4	50,042,880	493,818	101.3
D	3	2	61,273,088	612,488	100.0
E	3	4	28,986,368	308,318	94.0

計算量はこのコードで計算しました。あくまで「疑似計算量」なので正確な値とは異なると思います。パラメーター数はKerasのmodel.summary()のパラメーター数を見ました。ほとんど計算量はパラメーター数と比例しているのがわかります。

def calc_single_conv_layer(kernel, width, input_ch, output_ch):
    return kernel**2 * width ** 2 * input_ch * output_ch

def calc_single_normal_layer(kernel, width, input_ch):
    return kernel**2 * width ** 2 * input_ch

def calc_single_module(kernel, width, input_ch, output_ch):
    x = calc_single_conv_layer(kernel, width, input_ch, output_ch) # Conv
    x += calc_single_normal_layer(1, width, input_ch) # BatchNorm（簡易計算）
    x += calc_single_normal_layer(1, width, input_ch) # ReLU
    return x

def calc_single_module_with_bottleneck(mode, kernel, width, input_ch, output_ch, alpha):
    if mode == 1:
        return calc_single_module(kernel, width, input_ch, output_ch)
    elif mode == 2:
        # Bottleneck -> Conv
        x = calc_single_module(1, width, input_ch, output_ch//alpha)
        x += calc_single_module(kernel, width, output_ch//alpha, output_ch)
        return x
    elif mode == 3:
        # Bottlneck -> Conv -> Bottleneck
        x = calc_single_module(1, width, input_ch, output_ch//alpha)
        x += calc_single_module(kernel, width, output_ch//alpha, output_ch//alpha)
        x += calc_single_module(1, width, output_ch//alpha, output_ch)
        return x

def calc_inception_module(width, input_ch, output_ch, alpha, mode):
    assert output_ch % (alpha * 8) == 0
    assert mode >= 1 and mode <= 3 and type(mode) is int
    # conv1
    x = calc_single_module(1, width, input_ch, output_ch//4)
    # conv3, 5
    x += calc_single_module_with_bottleneck(mode, 3, width, input_ch, output_ch//2, alpha)
    x += calc_single_module_with_bottleneck(mode, 5, width, input_ch, output_ch//8, alpha)
    # pool
    x += calc_single_normal_layer(3, width, input_ch)
    x += calc_single_module(1, width, input_ch, output_ch//8)
    return x

def calc_model_flops(mode, alpha):
    x = calc_inception_module(32, 3, 96, alpha, mode)
    # pool 32->16
    x += calc_single_normal_layer(2, 32, 96)
    x += calc_inception_module(16, 96, 256, alpha, mode)
    # pool 16 -> 8
    x += calc_single_normal_layer(2, 16, 256)
    x += calc_inception_module(8, 256, 384, alpha, mode)
    x += calc_inception_module(8, 384, 384, alpha, mode)
    x += calc_inception_module(8, 384, 256, alpha, mode)
    return x

m1 = calc_model_flops(1, 2)
m2 = calc_model_flops(2, 2)
m3 = calc_model_flops(2, 4)
m4 = calc_model_flops(3, 2)
m5 = calc_model_flops(3, 4)

print(f"{m1:,d}\t{m2:,d}\t{m3:,d}\t{m4:,d}\t{m5:,d}")

ただし、残念なことにボトルネック層の使って計算量を減らしても、訓練時間が縮むかというと必ずしもそうではないです。以下、KerasとPyTorchの2つのフレームワークを使って確かめてみますが、計算量が1/5ぐらいになっても1エポックあたりの訓練時間がかえって増えてしまっているケースもあります。そこは割り引いて読む必要はあります。

結果

FW	パターン	mode	α	val_acc	(sd)	s / epoch	(sd)
Keras	A	1	-	88.67%	0.21%	76.01	1.46
Keras	B	2	2	90.48%	0.18%	79.77	2.56
Keras	C	2	4	90.33%	0.24%	68.30	0.81
Keras	D	3	2	90.88%	0.28%	94.76	2.08
Keras	E	3	4	90.33%	0.21%	84.47	1.19
PyTorch	A	1	-	89.37%	0.14%	46.75	0.26
PyTorch	B	2	2	90.85%	0.13%	45.58	0.16
PyTorch	C	2	4	90.21%	0.29%	40.39	0.36
PyTorch	D	3	2	90.83%	0.09%	48.53	0.40
PyTorch	E	3	4	89.90%	0.22%	41.47	0.21

左から、フレームワーク、パターン、ボトルネックのモード（1はボトルネックなし、2はBottleneck→Conv、3はBottleneck→Conv→Bottleneck）、ボトルネックのα、最大Validation Accuracyの中央値とその標準偏差、1エポック目以外のエポックあたりの訓練時間の中央値とその標準偏差です。KerasのバックエンドはTensorFlowを使っています。

訓練条件は以下の通りです。各パターンについて5回試行しました。DataAugmentationとしてよく用いられる水平反転＋ランダムクロップ、L2正則化として1e-3を付与しました。オプティマイザーはAdam、初期学習率は1e-3で、全てのパターンにつき100エポック訓練させ、全体の50%と85%で学習率を1/10にそれぞれ減らしています。全てGoogle ColabのGPUで訓練させました。

この結果はわかることは以下の通りです。

• PyTorchはボトルネックによる時間削減効果が出ているが、Kerasは効果が出づらい。
  PyTorchでは、パターンDという例外はあるが、ボトルネックなしのパターンAはボトルネックありの他のパターンと比べて計算時間が長い。Kerasでは処理時間はA<C<Eとなっているので、α=2ではボトルネックを使っても処理時間は短くならない。KerasではパターンAより速くなったのがパターンCしかない。なぜ遅くなったのかは不明だが、ボトルネックによる計算量削減よりも、モデルが深くなったことへのオーバーヘッドのほうが大きかったのでは。PyTorchでは比較的素直な結果が出てきたので、Kerasが遅いと言われるのはこのへんが絡んでいるかもしれない。

• どのフレームワーク・モードでも、αを大きくすると必ず処理時間は短くなる。しかし、精度とトレードオフにはなる。
  もしモデルの深さが同じなら、ボトルネック層での次元削減のパラメーターαを大きくすると、必ず処理時間は短くなっている。これは明確にボトルネックに層よる計算量削減効果が出ている証。ただし、理論的に計算すると、例えばB→Cで計算量が半分ぐらいになっているのに、処理時間は半分になってはいない。なぜなら、畳み込みの計算以外のオーバーヘッドがあるから。計算量を半分近くにしたのに、結果的には1～2割ぐらいしか短縮できなかったというのはちょっと悲しい。
  ただし、αを大きくし圧縮率を上げ高速化することと精度はトレードオフになる。αが2なら、ReLUの非線形性によりモード2,3ともにモード1より良くなっているが、αが4だとモード1とどっこいどっこまで落ちているケースもある（PyTorchのEなど）。したがってαを極端に大きくするのはよくない。

ただし、訓練時間は目に見えた変化なくても、モデルのパラメーターは減るという明確なメリットはあるので、アプリケーションにしたときにパラメーターの容量を減らしたいのなら大いに使うべきだと思います。なにかと係数の容量がでかくなりがちなので。

まとめ

• ボトルネック層を使うと理論的な計算量はガクッと落ちる。ただし実際の訓練時間が減るかどうかは、モデルが深くなることによるオーバーヘッドがあるので、フレームワークごとにまちまちで高速化しないこともある。
• ボトルネック層のチャンネル圧縮の比率（α）を上げると、明確に高速化するが精度とトレードオフになる。常識的な比率（2とか）なら、ReLUの非線形性により高速化と精度の向上を両立できる。
• パラメーターの容量はボトルネックを使うと明確に軽くなる。

計算の重そうな畳み込みはどんどんボトルネックを入れてみよう！ということで以上です。