TensorFlow2.1のアップデート情報: Keras と TPU と再現可能性の話

これは TensorFlow 2.0 AdventCalendar 2019 最終日の記事です。

今年は TensorFlow 2.0 のリリースなど、TensorFlow界隈に大きな変化のある年でした。
このアドベントカレンダーも、 TensorFlow 2.0 についてのものですが、実はすでに 2.1 の開発が進んでおり、本記事の執筆時点では 2.1.0-rc2 となっています。

そこで、TensorFlow 2.1 について紹介して、 TensorFlow 2.0 AdventCalendar を締めたいと思います¹。

TensorFlow 2.1 における主な変更点

TensorFlow 2.1 でどのような変更が加わるかは、リリースノート で確認することができます。個人的に気になったのは以下の通りです。

1. pip パッケージ名の変更
   今までpipでGPUを使うには tensorflow-gpu を使う必要がありましたが、今後は tensorflow pip パッケージにGPUサポートが含まれるようになります。GPUの無い環境でも動くので、これまでのようにインストール先によってtensorflowとtensorflow-gpuを使い分ける必要がなくなります。

2. Keras の分散処理、TPU サポートの強化
   おそらくTensorFlow1.xの時代は、Keras よりも Estimator が推奨されていたという背景もあり、分散処理やTPUへの対応状況は、Estimatorのほうが進んでいる状況でした。今回のリリースで、KerasのTPUサポートがかなり強化されるようです。

3. Mixed Precision API
   こちらのブログでも少し言及しましたが、普段float32で演算している部分をfloat16（もしくはbfloat16）にすることで、計算速度の向上が見込めますが、むやみにfloat16にしてしまうと、精度の低下の原因となりえます。Mixed Precision API を使うと、float32とfloat16を適切に組み合わせることで、精度の低下を抑えつつ、高速化することができます。

4. 再現可能性の強化
   GPU上でTensorFlowを動かすと、シードを固定してもちょっとだけ結果が違うような状況があり、正直困っていました。環境変数 TF_DETERMINISTIC_OPS を "true" もしくは "1" と指定することで、cuDNN の convolution や max-pooling の挙動（しいては Keras の ConvD や MaxPoolD レイヤーの挙動）が再現可能になります。

Keras の分散処理、TPUサポートの強化

（実験的にですが） すべての Keras モデル（Sequential, Functional API、 サブクラス）の compile、fit、evaluatite、predict が TPU をサポートしました。
TPUを含む分散訓練の書き方については公式のチュートリアルがわかりやすいと思いますが、以下のようなコードになります。strategyの部分を変えると、TPU以外の分散方法にかえられます。

# まずは TPUStrategy を構築する
tpu_address = "grpc://" + os.environ["COLAB_TPU_ADDR"] # Colaboratory の場合
tpu_cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu_address)
tf.config.experimental_connect_to_cluster(tpu_cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(tpu_cluster_resolver)   
strategy = tf.distribute.experimental.TPUStrategy(tpu_cluster_resolver)

...

# TPUStrategy の scope の中でモデルを構築する
with strategy.scope():
    model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))])

model.compile(loss='mse', optimizer='sgd')
model.fit(dataset, epochs=2)
model.evaluate(dataset)

今回のリリースでは、Keras + TPUでの可変バッチサイズのサポートなど、細かい改善が多数なされています。また、tf.data についても、分散処理用の改善がはいり、パフォーマンスが改善するようで、分散処理に力を入れているのがわかります。

Mixed Precision API

Mixed Precision API についての説明は、公式のチュートリアルと、この機能のベースとなっている Mixed Precision Training という論文が参考になるかと思います。

近年、深層学習専用のアクセラレータがいくつか発表されていますが、例えば Google の TPU は、基本的に bfloat16 で演算をするように設計されています。同様に NVIDIA の Tensor Core は float16 の演算をサポートしています。そのため、これらのアクセラレータを使って訓練の高速化をするには、bfloat16 や float16 を使う必要があります。Mixed Precision API を使うと、 float16 や bfloat16 と float32 をうまく使いこなすことで、精度への影響を抑えつつ、訓練を高速化できます。

注意が必要なのは、劇的な高速化にはアクセラレータが必要である点です。例えば、Tensor Core は V100 や RTXシリーズなど、最近のGPUでなければ搭載されていません。例えば Colaboratoryで使えるGPUは現在はP100なので、残念ながら Tensor Core は搭載されていません（古いGPUでもメモリ消費量は減るので、メリットがないわけではありません）。

Mixed Precision API を使うには、weightを定義する際にdtypeを明示的に指定するのをやめ、プログラムの冒頭部分でポリシーを指定します。例えば以下の例だと、"mixed_float16" というポリシーに従うことになります。

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_policy(policy)

この "mixed_float16" というポリシーでは、変数はfloat32で保持しておきますが、計算自体はfloat16で行います。float32は勾配のアップデートのときに使い、数値計算の安定性を担保します。また、僕の理解が正しければ、行列演算では要素毎に積を取った結果の和をとりますが、その和の部分はfloat32で行います（Tensor Core でサポートされている機能）。

現在のポリシーがどのように float32 と float16 を使い分けているかは、以下のようにして確認できます。

print('Compute dtype: %s' % policy.compute_dtype)
print('Variable dtype: %s' % policy.variable_dtype)

Compute dtype: float16
Variable dtype: float32

ちなみに、損失関数に渡す値がfloat16だと、勾配計算のときにオーバーフロー/アンダーフローの問題になりがちなので、以下のように最終層はfloat32にしておくのが定石のようです。

x = layers.Dense(10, name='dense_logits')(x)
outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)

ちなみに、NVIDIA関係では TensorRT6.0 がサポートされ、デフォルトで使用されるようになったようです。

まとめと言う名の雑感

読んでいただいてわかるとおり、今回はTPU周りがメインのリリースかと思いますが、個人的には再現可能性の強化がめっちゃ気になるので、今後ちゃんと調べてみようと思います。というか、GPUつかうと結果を再現できない問題、エンジニア視点だと結構重要で、社内で幾度となく質問受けては落胆されていたので。。。

1. ちなみに、Colaboratry はすでに 2.1.0-rc1 になっているようです。 ↩