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

これは [TensorFlow 2.0 AdventCalendar 2019](https://qiita.com/advent-calendar/2019/tensorflow2) 最終日の記事です。

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

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

[^1]: ちなみに、Colaboratry はすでに 2.1.0-rc1 になっているようです。

## TensorFlow 2.1 における主な変更点

TensorFlow 2.1 でどのような変更が加わるかは、[リリースノート](https://github.com/tensorflow/tensorflow/blob/v2.1.0-rc2/RELEASE.md) で確認することができます。個人的に気になったのは以下の通りです。

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
   こちらの[ブログ](http://blog.brainpad.co.jp/entry/2019/11/01/151250)でも少し言及しましたが、普段float32で演算している部分をfloat16（もしくはbfloat16）にすることで、計算速度の向上が見込めますが、むやみにfloat16にしてしまうと、精度の低下の原因となりえます。Mixed Precision API を使うと、float32とfloat16を適切に組み合わせることで、精度の低下を抑えつつ、高速化することができます。

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

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

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

```python
# まずは 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 についての説明は、公式の[チュートリアル](https://www.tensorflow.org/guide/keras/mixed_precision)と、この機能のベースとなっている [Mixed Precision Training](https://arxiv.org/abs/1710.03740) という論文が参考になるかと思います。

近年、深層学習専用のアクセラレータがいくつか発表されていますが、例えば 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" というポリシーに従うことになります。

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

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

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

```python
print('Compute dtype: %s' % policy.compute_dtype)
print('Variable dtype: %s' % policy.variable_dtype)
```
```
Compute dtype: float16
Variable dtype: float32
```

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

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

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

## まとめと言う名の雑感

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