ニューラルネットの新しい正規化手法 Group Normalization の高速な実装と学習実験

今年 1 月に ALBERT に入社した清水です。深層学習プログラマとして自社プロダクト開発をしております。このブログを書くのは始めてなのですが、今日はちょっとプログラミング寄りの記事を。残暑厳しい折ですが、実装の詳細にまで立ち入りつつアツく Yuxin Wu および Kaiming He の考案した手法 Group Normalization（ECCV 2018 採択済）について語ります。Kaiming He 氏は ResNet を筆頭に優れた convolutional neural networks (CNN) の設計で知られていることもあり、みなさんも注目している手法ではないでしょうか？

Chainer v5.0.0b4 を用いて、Chainer に入っている実装と弊社の独自実装での速度比較、また Batch Normalization との速度・精度比較を行いました。その上で、速度差が生じる原因の調査や CUDA を使った高速化の工夫について詳しく記します。ソースコードは MIT ライセンスで GitHub に公開していますので、ぜひご覧ください。

Group Normalization って？

Group Normalization の発想はシンプルです。早速ですが、Group Normalization を端的に表す図（論文 Figure 2 より引用）を見てみましょう。

(N, C, HW) の 3 次元からなる多次元配列が出てきましたが、これは CNN の中間層の出力だと考えると想像しやすいかと思います。バッチサイズ N, チャンネル数 C, 画像の縦幅 H, 画像の横幅 W です。

図に示された Batch Normalization, Layer Normalization, Instance Normalization およびこの記事の本題 Group Normalization の 4 つの手法は、いずれも青色で示された領域ごとに算出される平均・分散によって入力を正規化 (normalize) します。Batch Normalization が各チャンネルごとに平均・分散を求めることは有名ですね。精度・収束性を劇的に向上させる Batch Normalization は今や CNN のデファクトスタンダードですが、しかし

• 画像の解像度が大きくてメモリが不足するなどの理由でバッチサイズが小さくなる場合に、平均・分散の推定が不安定になり学習ができなくなる
• 複数 GPU にまたがって平均・分散を推定することで実質的なバッチサイズを増やすことは可能だが、高価なハードウェアが必要になる上に実装や最適化が複雑
• ビデオの隣接フレームといった相関がある画像をミニバッチとして入力する場合も平均・分散の推定が不安定になる
• 学習時に平均・分散の移動平均を覚えておいて推論時に用いるといった処理が煩雑
• Finetune 時に移動平均をどう扱うべきかよくわからない

といった難点も併せ持っており、いつでも使えるわけではありません。このためミニバッチに依存しない正規化手法が待ち望まれていました。

そのような手法として、全チャンネルにまたがって平均・分散を取る Layer Normalization と、各チャンネル独立に画像の縦横方向についてのみ平均・分散を取る Instance Normalization が考案されました。とはいえ十分にバッチサイズが確保されている条件下での精度は Batch Normalization に比べてかなり劣っており、主流とはなっていません。

そこで登場したのが Group Normalization です。チャンネルを G 個にグルーピングして Layer Normalization と Instance Normalization の中間的な処理を行うことで、画像分類などのタスクで Batch Normalization に匹敵する精度を実現しました。グループ数 G を 32 にした場合がベストだったと論文に述べられていますが、それほど G の値に対して敏感ではないようです。Group Normalization の論文では、Instance Normalization には複数のチャンネルの組み合わせによって表現される特徴を歪めてしまう問題があると考察されています。その対極になるのが Layer Normalization ですが、こちらは大域的すぎて、特に重要ではないチャンネルが画像全体にわたって高い値を示した場合に他の全てのチャンネルが抑制されてしまいます。中間的な Group Normalization は良いとこ取りをできるというわけです。

なんだか素晴らしそうですね。Chainer では v5.0.0b3 から Group Normalization がサポートされているのでお手軽に使えます。しかし、本当に Batch Normalization をドロップインで置き換えて精度低下は起きないのでしょうか？ Batch Normalization と同等の速度で動作するのでしょうか？ この疑問を検証します。

結論から言えば、Batch Normalization を単に Group Normalization に置き換えるだけでは精度がかなり落ちてしまいました。なので Group Normalization を使う場合は精度の確認やパラメータチューニングをきちんとやるべきでしょう。また、Chainer v5.0.0b3 に追加された Group Normalization の実装はあまり効率的ではなく、CNN 全体の実行速度を大きく下げてしまうことがわかりました。この原因や、より効率のいい実装方法についても詳述します。

実験タスクおよび学習条件

ResNet-50 ネットワークを用いて、オリジナルの Batch Normalization を使うバージョンと、グループ数 G=32 の Group Normalization を用いるよう変更したバージョンでそれぞれ画像分類の学習を行いました。用いたデータセットは Oxford VGG Pets データセットに含まれる 12 品種の猫の画像です。データセットはなんでもよかったのですが

• CIFAR-10 もしくは CIFAR-100 だと実務で扱うデータと比べて画素数が少ないため、計算時間の測定条件としては不適切
• ImageNet のような大規模データセットの学習には時間がかかる
• 猫かわいい

という理由で決めました。

各品種ごとにおよそ 200 枚ずつのデータがあります。このうち 8 割を訓練用に、2 割を検証用に用いました。

学習に用いた設定は

• バッチサイズ: 32
• 画素数: 短辺が 256px になるようバイリニア補間でリサイズ後、224px × 224px の正方形にクロップ
• Augmentation: ランダム位置でのクロップおよび、ランダムな左右反転
• Optimizer: MomentumSGD, momentum=0.9
• Weight decay: 0.0005
• エポック数: 300
• 学習率: 0.1 から始めて、100 エポックごとに 1/10 倍

です。確実な実験を期するため、疑似乱数生成器のシードを固定することでデータセットのサンプリング順序・重みの初期化・augmentation の挙動が実験ごとに変わらないようにしています。ただし cuDNN を使っているため、アトミック加算などの誤差の出方が GPU 内でのスレッドスケジューリングに依存することによる、数値計算的な非決定性は存在します。

実験を行った計算機環境は

• ハードウェア
  • CPU: Intel Core i5-3550
  • GPU: ZOTAC GeForce GTX 1080 Ti AMP Edition
  • SSD: Crucial MX200 500GB
  • RAM: DDR3-1600 SDRAM 4GiB × 4
• ソフトウェア
  • Gentoo Linux amd64 (kernel 4.18.5)
  • Python 3.6.6
  • NVIDIA driver 396.54
  • CUDA 9.2 / cuDNN 7.1
  • GCC 7.3.0

となります。CPU がやや古めですが Pillow-SIMD を使うことで CPU が律速するという事態は回避できました。

Chainer v4.0.0 からは MultithreadIterator が追加されてデータの並列読み込みが簡単になったのですが、試してみたところ MultiprocessIterator でないと GPU がフルに動いてくれませんでした。プロセスの実行順序に依存しない augmentation を実現するためにデータセットの各要素ごとに疑似乱数生成器の内部状態を持たせる設計を採用し、実装にあたっては全ワーカープロセスからアクセスできる shared memory を用いました。また（ほとんど趣味ですが）疑似乱数アルゴリズムの勉強も兼ねて、メモリ消費の少ない xoshiro128+ 疑似乱数生成器を Python に移植して使うことでフットプリントを低減してみました。

学習結果

正規化のレイヤーを

• Batch Normalization
• Group Normalization (Chainer の実装)
• Group Normalization（ALBERT の実装 1 つ目）
• Group Normalization（ALBERT の実装 2 つ目、高速版）

のそれぞれに入れ替え、ResNet-50 の学習を行った結果を以下に示します。先述したように Group Normalization におけるグループ数 G は 32 に固定しています。検証は毎エポックごとに中心クロップ画像に対して行いました。

Batch Normalization の圧勝ですね……。学習率・weight decay・グループ数などの調整を行うと違った結果にはなるでしょう。

Group Normalization を使う 3 つのバージョンは実装の違いだけですので、当然ながらほぼ同じ結果になっています。ただ、ALBERT1 と ALBERT2 の実装は浮動小数点演算に関しては全く同一で、cuDNN による非決定性以外に違いの出る要因はないはずなのですが、にもかかわらず train accuracy にかなりの差が出たのは印象的でした。

次に速度の比較です。ここで挙げる実行時間は学習・検証全体で測定したものです。Normalization 単体のマイクロベンチマークをすればさらに差は開くでしょう。いずれにせよ Batch Normalization は cuDNN で実装されているので、勝つのはなかなか困難だと予想されます。

Batch Normalization では 52 分で終わった学習が、Chainer の Group Normalization 実装では 1.66 倍の 86 分かかってしまっていました。一方で ALBERT2 版では 1.06 倍と、かなり近い実行時間にまで最適化できています。

学習ログデータとプロットに用いたプログラムも GitHub に公開しています。

Chainer v5.0.0b3 の Group Normalization 実装ではどうして時間がかかっている？

Chainer v5.0.0b3 (v5.0.0b4 地点でも同じ) の Group Normalization のソースコードは、入力をうまく reshape することで Batch Normalization を行う関数を利用する巧妙な実装となっています。正規化処理本体は cuDNN などのバックエンドに備わった Batch Normalization ルーチンに任せてしまい、後は各チャンネルにスケール \gamma とバイアス \beta を適用するだけというわけです。Batch Normalization の呼び出しと \gamma によるスケーリングのそれぞれが backward のためのメモリを確保するためにメモリ消費が倍になる弱点はあるものの、一見したところ高速に動作しそうです。（なお適切な再計算をすればメモリ確保はほぼなくせますが、この記事では扱いません。既に詳しい実装や論文が公開されています。）

そこで NVIDIA Visual Profiler で調べてみたところ、特別遅い箇所は確かに存在しないものの、GPU に発行される多数の細かい処理（CUDA カーネル）が積もり積もることで時間が消費されているとわかりました。さらに

nvidia-smi

コードを見つつ、具体的に何回の CUDA カーネル発行が起こっているか確認してみましょう。

chainer/functions/normalization/group_normalization.py

    # By doing this reshaping, calling batch_normalization function becomes
    # equivalent to Group Normalization.
    # And redundant dimension is added in order to utilize ideep64/cuDNN.
    x = reshape.reshape(x, (1, batch_size * groups, -1, 1))

    with cuda.get_device_from_array(x.array):
        dummy_gamma = xp.ones(batch_size * groups).astype(xp.float32)
        dummy_beta = xp.zeros(batch_size * groups).astype(xp.float32)

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        x = batch_normalization.batch_normalization(
            x, dummy_gamma, dummy_beta, eps=eps)

    x = reshape.reshape(x, original_shape)

    target_shape = [1, channels] + [1] * (x.ndim - 2)
    gamma_broadcast = broadcast.broadcast_to(
        reshape.reshape(gamma, target_shape), x.shape)
    beta_broadcast = broadcast.broadcast_to(
        reshape.reshape(beta, target_shape), x.shape)

    return x * gamma_broadcast + beta_broadcast

Forward 時、backward 時それぞれについて調べてみると以下の処理が GPU で行われていることがわかります。

• Forward 時
  • 59行目:

    chainer.functions.reshape

    内部でのメモリコピー
  • 62行目:

    xp.ones

    によるメモリを 1 で埋める処理
  • 62行目:

    cupy.ndarray.astype

    による型変換
  • 63行目:

    xp.zeros

    によるメモリを 1 で埋める処理
  • 63行目:

    cupy.ndarray.astype

    による型変換
  • 67行目:

    chainer.functions.batch_normalization

    内部での、ダミーの rolling average を 0 で埋める処理
  • 67行目:

    chainer.functions.batch_normalization

    内部での、ダミーの rolling variance を 1 で埋める処理
  • 67行目: cuDNN の Batch Normalization foward 処理呼び出し
  • 70行目:

    chainer.functions.reshape

    内部でのメモリコピー
  • 74行目:

    chainer.functions.reshape

    内部でのメモリコピー
  • 73行目:

    chainer.functions.broadcast_to

    内部でのメモリコピー
  • 76行目:

    chainer.functions.reshape

    内部でのメモリコピー
  • 75行目:

    chainer.functions.broadcast_to

    内部でのメモリコピー
  • 78行目:

    *

    による乗算
  • 78行目:

    +

    による加算
• Backward 時
  • 78行目:

    +

    の backprop に伴うメモリコピー（2 回）
  • 78行目:

    *

    の backprop に伴う乗算（2 回）
  • 75行目:

    chainer.functions.broadcast_to

    の backprop に伴う総和計算
  • 76行目:

    chainer.functions.reshape

    の backprop に伴うメモリコピー
  • 73行目:

    chainer.functions.broadcast_to

    の backprop に伴う総和計算
  • 74行目:

    chainer.functions.reshape

    の backprop に伴うメモリコピー
  • 70行目:

    chainer.functions.reshape

    の backprop に伴うメモリコピー
  • 67行目: cuDNN の Batch Normalization backward 処理呼び出し
  • 59行目:

    chainer.functions.reshape

    の backprop に伴うメモリコピー

どうしたって必要な cuDNN 呼び出し以外に、forward で 14 回、backward で 10 回の CUDA カーネル呼び出しが起こっていたのです。Chainer の自動微分は便利ですが細かいカーネルが増えがちです。特にネットワーク内部で何度も使われる関数を自動微分に頼って書くと、全体の実行時間に大きな影響を与えうることがわかります。

この実装の場合、cuDNN 呼び出し以外の本質的な処理は 78 行目だけですので、がんばって色々修正すれば forward / backward ともに CUDA カーネル 1 回ずつにまで低減することは可能です。深層学習フレームワークは急速に進歩していますので、近い将来には自力でがんばらなくても自動的に最適化されるようになるでしょう。

しかし、その方針ではどうしても cuDNN の Batch Normalization に勝てませんし、メモリ確保も減らせません。なので cuDNN を使わずフルスクラッチで実装することにしました。

Group Normalization の誤差逆伝搬

フルスクラッチで実装するなら、まずすべきことは微分です。入力となる多次元配列 x の shape を (N,C,H,W) であるとし、添字ベクトル i を (i_N,i_C,i_H,i_W) によって表すことにします。チャンネル数 C がグループ数 G で割り切れる場合のみを考えることにして、グループあたりのチャンネル数 C/G を C_G と表記します。さらに、チャンネル i_C が何番目のグループに属するかを i_G=\lfloor i_C/C_G \rfloor で表記します。

正規化済みデータ \hat{x} を平均 \mu と分散 \sigma^2 を用いて \hat{x}_i=(x_i-\mu_{i_N,i_G})/\sigma_{i_N,i_G} によって定義すると、出力 y は y_i=\gamma_{i_C}\hat{x}_i+\beta_{i_C} となります。

まず \mu と \sigma^2 でロス l を微分します。添字集合 S_{i_N,i_G} を \{j\mid j_N=i_N,\; j_G=i_G\} によって定義しておきます。微分の結果を見ると、この添字集合 S_{i_N,i_G} ごとに総和計算が走ることがわかります。

これらがわかれば入力 x に関する勾配を求めることができます。

スケール \gamma とバイアス \beta に関する勾配は単純です。こちらでは添字集合 S_{i_C}=\{j \mid j_C=i_C\} ごとに総和計算を走らせることになります。

Group Normalization ALBERT1 版の実装

CuPy では

compile_with_cache

chainer.backends.cuda.elementwise

chainer.backends.cuda.reduce

chainer.backends.cuda.raw

Forward 時には、shared memory を用いて GPU のスレッド間で値を集約する標準的なアルゴリズムを用いて、まず各グループ S_{i_N,i_G} 内での総和 \sum_{i\in S_{i_N,i_G}}x_i および二乗和 \sum_{i\in S_{i_N,i_G}}x_i^2 を同時に求めました。そこから平均 \mu および標準偏差の逆数 \sigma^{-1} を計算できます。（分散 \sigma^2 が \Big(\sum_{i\in S_{i_N,i_G}}x_i^2\Big)/C_GHW-\mu^2+\varepsilon で求まることを用いています。）あとは普通に \hat{x}=(x-\mu)/\sigma および y=\gamma\hat{x}+\beta を

elementwise

Backward の実装はより面倒です。誤差逆伝搬の式展開からわかるように S_{i_N,i_G} と S_{i_C} の両方について総和計算が必要になるからです。一旦 HW 方向についての総和

を求めておいて、そこから C_G 方向および N 方向の総和をそれぞれ取って

と計算するのがよさそうです。そのまま実装すると総和計算のために 3 回の CUDA カーネル呼び出しが必要になりますが、A_{i_N,i_C} および \hat{A}_{i_N,i_C} を求めるカーネル内に C_G 方向の総和を取る処理も詰め込んでしまうことで、カーネル呼び出しを 1 つ削減しています。

なお今回の実験では問題にならなかったものの、分散 \sigma^2 の計算は数値的に不安定です。このため double 型で総和を取ったり、オンライン / 並列アルゴリズムで残差二乗和を計算したり、平均と分散を 2 パスで計算したりして数値精度を改善することで、出力値を安定させられる可能性はあります。これらの計算方法については Algorithms for calculating variance が詳しいです。また総和計算の速度面では、shared memory を使わず shfl 命令を使ったり、段階的なカーネル実行によって並列度を上げたりといった最適化が考えられます。Faster Parallel Reductions on Kepler などを参照ください。

Group Normalization ALBERT2 版の実装

ALBERT1 版の実行時間は満足いくものではなかったので、また NVIDIA Visual Profiler のお世話になりました。どうやら forward / backward 時の

elementwise

elementwise

落とし穴はインデックスの計算にありました。Numpy の broadcasting rule にしたがえば y=(x-\mu)\sigma^{-1}\gamma+\beta を素直なソースコードで計算できるのですが、ここで x の shape は (N, G, C_G, HW)、\mu, \sigma^{-1} の shape は (N, G, 1, 1)、\gamma, \beta の shape は (1, G, C_G, HW) と、バラバラになっています。この場合に i 番目のスレッドがアクセスすべきメモリアドレスを求めるためには、整数除算を用いて i_{HW}=i % HW,\; i_{C_G}=i / HW % C_G,\; i_{G}=i / HW / C_G % G,\; i_{N}=i / HW / C_G / G というふうに各軸のインデックスを計算する必要があるのです。

整数除算は CPU や GPU にとって非常に時間のかかる処理です。なんとしても消し去らねばなりません。

現実的には多くの場合チャンネル数やバッチサイズは 2 の冪なので、除算をビットシフトに置き換えることができます。そうでなくとも、コンパイル時に除数がわかっているなら乗算とビットシフトの組み合わせなどに変形可能です。（“Hacker’s Delight” や日本語訳『ハッカーのたのしみ』の 10 章に詳しい説明があります。）

したがい、コードの中に

n_elems_per_channel

n_channels_per_group

groups

elementwise

group_normalization_alb2_func.py

        gn_fwd_norm_code = string.Template('''
            unsigned int u = i;
            unsigned int tmp = u / ${n_elems_per_channel};
            unsigned int channel_idx = tmp % ${n_channels_per_group};
            tmp /= ${n_channels_per_group};
            unsigned int group_idx = tmp % ${groups};
            unsigned int batch_idx = tmp / ${groups};
            unsigned int group_norm_idx =
                batch_idx * ${groups} +
                group_idx;
            unsigned int batch_norm_idx =
                group_idx * ${n_channels_per_group} +
                channel_idx;
            T v_x_hat = (x[u] - mu[group_norm_idx]) * inv_std[group_norm_idx];
            x_hat[u] = v_x_hat;
            y[u] = v_x_hat * gamma[batch_norm_idx] + beta[batch_norm_idx];
        ''').substitute(n_elems_per_channel=n_elems_per_channel,
                        n_channels_per_group=n_channels_per_group,
                        groups=groups)
        gn_fwd_norm_name = 'gn_fwd_norm_{}_{}_{}'.format(
                    n_elems_per_channel, n_channels_per_group, groups)
        gn_fwd_norm_kern = cuda.elementwise(
            'raw T x, raw T mu, raw T inv_std, raw T gamma, raw T beta',
            'raw T x_hat, raw T y',
            gn_fwd_norm_code, gn_fwd_norm_name)

        x_hat, y = gn_fwd_norm_kern(x, mu, inv_std, gamma, beta,
                                    size=x.size)

この修正を加えることで、完璧とは言えないものの、メモリアクセス律速と思われる Batch Normalization に近い速度にまで持っていくことができました。CuPy 側でも軸のサイズが 2 の冪の場合の高速化の試みはありますが、執筆地点ではまだマージされていないようです。

なんにせよ、かなり強引な実装なので Chainer へのプルリクエストは現地点では考えていません。特に backward 側は勢いでベンチマークもせずに複雑な実装をしてしまった感は否めず、素直に総和計算を 3 回行う実装でもよかったかもしれないと感じています。CuPy 側の機能として

• 複数の値を同時に計算する reduction
• インデックス計算の自動的な最適化

がサポートされれば、かなりシンプルなコードに書き直せるんじゃないかなあと期待しているところです。

まとめ

• Group Normalization はミニバッチの取り方に依存することによる Batch Normalization の難点を克服し、かつ既存手法の良いとこ取りをすることで精度的にも比較的優れている
• しかし手元で Batch Normalization と同じ学習パラメータで実験したところ精度はかなり劣ったので、実戦投入する際には実データでの結果の確認やパラメータチューニングが必要になりそう
• Chainer v5.0.0b4 地点での Group Normalization の実装は細かい CUDA カーネルの呼び出しが多数発生しているために非効率
• フルスクラッチで Group Normalization を実装してみるとインデックス計算が律速するという落とし穴があり、そこを修正すると一気に高速化した

かなり高速化にこだわった内容でしたが、それが果たして重要な「イシュー」であるのか疑問を持たれた方もいるでしょう。しかしあくまで私見ですが、高速なプログラムは質のいいデータ分析を行うために極めて重要です。速ければ速いほど、様々な仮説を実験的に調べることでこれまで気づかなかった知見を見い出したり、細やかなパラメータチューニングによって精度を上げたりすることができるようになっていきます。ALBERT には最新の手法やオリジナルの手法を自ら Chainer などで実装できるアナリストが数多くいますが、やはり本格的な高速化となるとプログラマの領分になります。深層学習は今まさに発展著しい分野ですので、既存の高速なライブラリを使えないこともよくあります。

とはいえプログラマだけでデータ分析はできません。ビジネスでの実データの扱いや、データ分析手法の背景への深い理解など、学ぶべきことがたくさんあると痛感しています。深層学習ブーム以前からのベテランアナリストや、数学に強い理論派アナリスト、あるいはニューラルネットワークにいち早く目をつけて経験を蓄積した職人派アナリスト、さらにビジネスパートナーとのやり取りを円滑に調整してくださる営業やマネージャーなどなど、さまざまな人の力を借りて実践的なデータ分析に携わることができています。そもそもプログラマだとかアナリストだとか営業だとかの区別は曖昧で、個々人にそれぞれの強みがあるというべきなのでしょう。ALBERT は一緒に働くメンバーを募集しています！

私が携わっている ALBERT の画像分析支援サービスタクミノメ には Group Normalization がすでに組み込まれており、実際に分析業務で使われ始めています。顧客に最大の価値をもたらすこと・データ分析者にとっての使いやすさを大切にすること・シンプルかつ効率的な実装を心がけることで、ALBERT ならではの高水準なサービスを提供することができていると自負しております。