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

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 を締めたいと思います1

TensorFlow 2.1 における主な変更点

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

  1. pip パッケージ名の変更
    今までpipでGPUを使うには tensorflow-gpu を使う必要がありましたが、今後は tensorflow pip パッケージにGPUサポートが含まれるようになります。GPUの無い環境でも動くので、これまでのようにインストール先によってtensorflowtensorflow-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 の Conv*D や MaxPool*D レイヤーの挙動)が再現可能になります。

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 になっているようです。 

Why not register and get more from Qiita?
  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