Help us understand the problem. What is going on with this article?

Automatic Mixed Precision (AMP) でニューラルネットワークのトレーニングを高速化

エヌビディアの佐々木です。

ニューラルネットワークのトレーニングに Volta 及び Turing アーキテクチャの (つまりわりと新しい) GPU をご利用の方は、 Tensor コアによる混合精度演算をぜひ活用してください。Automatic Mixed Precision 機能を使えば既存のモデルを修正する必要がほぼなく、しかし場合によっては 3 倍くらい速くなるかもしれません。詳細はこちら → Automatic Mixed Precision for Deep Learning

と、これだけで終わるのも何なので以下、説明します。

混合精度演算

Mixed Precision Training という Baidu Research と NVIDIA による論文があります。この中では、従来ニューラルネットワークのモデルで一般的に利用されてきた 32 ビットの「単精度浮動小数点数 (FP32)」に代えて、半分の 16 ビット幅で表現される半精度浮動小数点数 (FP16) を使っても、いくつかの工夫を加えることでモデルの正確度をほぼ落とすことなく、トレーニングを高速化できることが示されています。

FP32 の代わりに FP16 でモデルが表現できるとなれば、当然必要なメモリサイズは半分になり、その分バッチサイズを大きくできますし、ネットワーク越しに重みをやりとりする際の帯域も節約できます。そして、それだけではないのです。先の論文の Introduction にこんな一節があります:

half-precision math throughput in recent GPUs is 2× to 8× higher than for single-precision.
最近の GPU では、FP16 の演算スループットが FP32 の 2 倍から 8 倍高い

そうなんです。この「2倍」は Pascal アーキテクチャで導入された、「FP32 演算器に FP16 の数値を 2 個詰め込んで一度に演算する」機能で、例えば Tesla P100 SXM2 では FP16 の演算性能 (21.2 TFLOPS) が FP32 (10.6 TFLOPS) の 2 倍になっています。そして、ここからがこの記事の本題ですが「FP32 の 8 倍」という大幅な性能アップを実現するのが、 Volta アーキテクチャで導入され、Turing アーキテクチャでさらに拡張された Tensor コア (てんさーこあ) です。

Tensor コア使ってますか?

Tensor コアは、ニューラルネットワークのトレーニング高速化を主な目的とした追加された行列演算ユニットで、Volta / Turing GPU の性能を引き出すために不可欠な要素です。例えば Tesla V100 SXM2 の「125 TFLOPS」は Tensor コアを使った FP16/FP32 の混合精度演算性能であり、FP32 の 15.7 TFLOPS とは 8 倍の差があります。そう、先ほど引用した "8× higher than for single-precision" とはこのことです。

image.png
出典: VOLTA AND TURING: ARCHITECTURE AND PERFORMANCE OPTIMIZATION - エヌビディア成瀬, GTC Japan 2018

Tensor コアは今のところ (2019 年 6 月現在) 下記の Volta / Turing アーキテクチャ GPU に搭載されています。

しかし、せっかくこれらの GPU を使っていても、Pascal アーキテクチャ以前のプログラムをそのまま動かすだけでは、おそらく Tensor コアは眠ったまま!です。もったいないお化けが出そうですね。どんどん Tensor コアを使っていきましょう。そのためのオススメ機能が、この記事で紹介する Automatic Mixed Precision (AMP) です。

Automatic Mixed Precision (AMP)

前述の通り Tensor コアは FP16 に対する演算を行いますから、既存のモデルで Tensor コアを活用するためには、FP32 で表現されている数値を FP16 に変更する必要があります。しかし、FP16 は FP32 と比較して表現できる値の範囲がかなり狭いため、単に FP32 を FP16 に置き換えるだけではモデルの正確度を維持できません。Mixed Precision Training の論文では、モデルを FP16 化しても正確度を落とさない方法として、次の 3 つの工夫が挙げられています。

  1. (モデルを FP16 化しつつも) 重みのマスターコピーは FP32 で保持
  2. ロス スケーリング (勾配の値がアンダーフローして 0 になってしまう避けるために、順伝播の結果得られた損失関数の値を数倍してから逆伝播を行う)
  3. FP16 の行列積を FP32 で積算

Automatic Mixed Precision (AMP) 機能は、このうち 1. と 2. を自動的に行ってくれます。そして、3. は Tensor コアの機能そのものです。

とりあえず試してみる

論より証拠、AMP の動作を試してみましょう。今のところ TensorFlow, PyTorch, MXNet が AMP に対応しています。NGC から新しめのコンテナイメージを持ってくるのが一番簡単です。例えば TensorFlow であれば、19.03 版以降は AMP に対応しています。
※ TensorFlow 1.14 から、master ブランチにも AMP が取り込まれました。

私は、会社の PC で NGC の TensorFlow イメージ 19.06 版を試してみることにしました。イメージ内の /workspace/nvidia-examples に様々なサンプルプログラムが含まれていますので、今回は ResNet50 のトレーニングを試してみます。なお、データをその場で生成するベンチマークモード (--mode=training_benchmark) で実行するため、学習データの準備は不要です。

まず、ベースラインとして、AMP を有効化せずに実行してみます。

$ docker run --runtime=nvidia --rm -it nvcr.io/nvidia/tensorflow:19.06-py3 /workspace/nvidia-examples/resnet50v1.5/main.py --mode=training_benchmark --batch_size=96 --num_iter=100 --iter_unit=batch --results_dir=/tmp
~中略~
:::NVLOGv0.2.3 resnet 1561803648.137060165 (benchmark_hooks.py:66) Average total_ips: 278.36461497043854

Average total_ips というのが、秒間処理画像数の平均です。AMP を有効にしない場合 278 枚前後でした。また、バッチサイズを変化させながら何度か試しましたが、この場合 96 ぐらいが限度でした。

では、AMP を有効にしてみます。TensorFlow の場合、環境変数 TF_ENABLE_AUTO_MIXED_PRECISION に 1 をセットしてやるだけです。コード修正は不要!

$ docker run --runtime=nvidia --rm -it -e TF_ENABLE_AUTO_MIXED_PRECISION=1  nvcr.io/nvidia/tensorflow:19.06-py3 /workspace/nvidia-examples/resnet50v1.5/main.py --mode=training_benchmark --batch_size=192 --num_iter=100 --iter_unit=batch --results_dir=/tmp
~中略~
:::NVLOGv0.2.3 resnet 1561804003.889541388 (benchmark_hooks.py:66) Average total_ips: 501.32636036112996

AMP を有効にした結果、バッチサイズは 192 まで拡大できて、秒間 500 毎前後処理できるようになりました。環境変数一つで 1.8 倍速くなるのであればなかなか「美味しい」と言えるのでは。

PyTorch では?

PyTorch の場合 NVIDIA が開発している APEX というライブラリが AMP の機能を提供しています。NGC の PyTorch コンテナイメージ (19.04 版以降) には APEX が含まれており AMP をすぐに利用できます。/workspace/examples/apex/imagenet/main_amp.py から関連する部分をちょっと抜き出してみましょう。コメントは私が説明のために付けたものです。

# APEX から AMP をインポートして
from apex import amp, optimizers

# AMP を初期化して
model, optimizer = amp.initialize(model, optimizer, opt_level=args.opt_level, keep_batchnorm_fp32=args.keep_batchnorm_fp32, loss_scale=args.loss_scale)

# ロススケーリングして逆伝播
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward() 

簡単ですね!

MxNet では?

MXNet でも NGC コンテナイメージの 19.03 以降で AMP に対応しており、また最近 (2019年7月末) リリースされた MXNet 1.5.0 からは AMP が本体に取り込まれています。設定方法は PyTorch とほぼ同様です。NGC の MxNet イメージ 内の /workspace/examples/gluon/image-classification/train_imagenet.py をちょっと見てみましょう。

# AMP をインポートして
import mxnet.amp as amp

# 初期化して
amp.init()
amp.init_trainer(trainer)

# ロススケーリングして逆伝播
with amp.scale_loss(Ls, trainer) as scaled_loss:
    ag.backward(scaled_loss)

やっぱり簡単ですね!

Tensor コアを活用する多くのサンプルプログラム

AMP は Tensor コアを使った混合精度演算を始める手軽な方法ですが、AMP を使わずに自分で混合精度演算のコードを書くことももちろん可能です。NVIDIA Deep Learning Examples for Volta Tensor Cores には Tensor コアの性能を引き出すサンプルプログラムがたくさん用意されています。こちらもぜひご覧ください。高速化の効果はモデルによりますが、例えば PyTorch による Transformer のトレーニング では 4 倍以上の高速化が達成されています。

[補足] Pascal 以前の GPU で実行したら?

AMP を組み込んだコードを Tensor コアを持たない Pascal アーキテクチャ以前の GPU で実行した場合、AMP 機能の呼び出しがエラーとなることはありませんが、例えば下記のようなメッセージが出力され、モデルの FP16 への自動変換は行われません。当然、学習速度が向上することもありません。

No (suitable) GPUs detected, skipping auto_mixed_precision graph optimizer

また、AMP とは別の話ですが、 NGC のコンテナイメージは Pascal アーキテクチャ以降の GPU を正式サポートしています。そのため Kepler や Maxwell 世代の GPU では実行できないイメージがある点にご留意ください。

まとめ

Tesla V100 や Tesla T4、TITAN RTX、GeForce RTX20 系などの Volta / Turing アーキテクチャ GPU をご利用の方、 Tensor コアが眠っていませんか?ぜひ、混合精度演算でニューラルネットワークのトレーニングを高速化しましょう。NGC のコンテナイメージと Automatic Mixed Precision (AMP) 機能で手軽に始めるのがお勧めです。

関連情報

Why do not you register as a user and use Qiita more conveniently?
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away
Comments
Sign up for free and join this conversation.
If you already have a Qiita account
Why do not you register as a user and use Qiita more conveniently?
You need to log in to use this function. Qiita can be used more conveniently after logging in.
You seem to be reading articles frequently this month. Qiita can be used more conveniently after logging in.
  1. We will deliver articles that match you
    By following users and tags, you can catch up information on technical fields that you are interested in as a whole
  2. you can read useful information later efficiently
    By "stocking" the articles you like, you can search right away