TensorFlow 2.0 Alpha : ガイド : TensorFlow で分散訓練 (翻訳/解説)
翻訳 : (株)クラスキャット セールスインフォメーション
作成日時 : 03/21/2019
* 本ページは、TensorFlow の本家サイトの TF 2.0 Alpha の以下のページを翻訳した上で適宜、補足説明したものです：
* サンプルコードの動作確認はしておりますが、必要な場合には適宜、追加改変しています。
* ご自由にリンクを張って頂いてかまいませんが、sales-info@classcat.com までご一報いただけると嬉しいです。
ガイド : TensorFlow で分散訓練
概要
tf.distribute.Strategy はマルチ GPU、マルチマシンや TPU に渡り訓練を分散するための TensorFlow API です。この API を使用して、ユーザは既存のモデルと訓練コードを最小限のコード変更で分散できます。tf.distribute.Strategy はこれらの主要なゴールに留意して設計されています : * 容易に使用できて研究者、ML 技術者, etc. を含むマルチ・ユーザセグメントをサポートします。* 良いパフォーマンスを独創的に提供します。* ストラテジー間の容易な切り替え。
tf.distribute.Strategy は TensorFlow の高位 API、tf.keras と tf.estimator で単に数行のコード変更で利用できます。それはまたカスタム訓練ループ (そして一般に TensorFlow を使用する任意の計算) を分散するために使用できる API も提供します。TensorFlow 2.0 では、ユーザはプログラムを eagerly に、あるいは tf.function を使用してグラフ内で実行できます。tf.distribute.Strategy はこれらの実行モードの両者をサポートすることを意図します。このガイドの時間の殆どで訓練について語るかも知れませんが、この API は異なるプラットフォーム上で評価と予測を分散するためにも使用できることに注意してください。
すぐに見るように、貴方のコードで tf.distribute.Strategy を使用するためには非常に少しの変更が必要とされます。これは strategy-aware とするために TensorFlow の基礎的なコンポーネントを変更したからです。これは変数、層、モデル、optimizer、メトリクス、要約そしてチェックポイントを含みます。
このガイドでは、様々なタイプのストラテジーについてとそれらを異なる状況でどのように使用できるかについて話します。
# Import TensorFlow from __future__ import absolute_import, division, print_function !pip install -q tensorflow==2.0.0-alpha0 import tensorflow as tf
ストラテジーのタイプ
tf.distribute.Strategy は異なる軸に沿った多くのユースケースをカバーすることを意図しています。これらの組み合わせの幾つかは現在サポートされていて他のものは将来的に追加されます。これらの軸の幾つかは : * 同期 vs 非同期訓練です : これらはデータ並列で訓練を分散する 2 つの一般的な方法です。同期訓練では、総てのワーカーは入力データの異なるスライスに渡り同期的に訓練して、各ステップで勾配を集めます。非同期訓練では、総てのワーカーは入力データに渡り独立的に訓練して変数を非同期的に更新します。典型的には同期訓練は all-reduce を通してそして非同期 (訓練) はパラメータサーバ・アーキテクチャを通してサポートされます。* ハードウェア・プラットフォーム: ユーザは一つのマシンかネットワーク内のマルチマシン (それぞれ 0 かそれ以上の GPU) 上のマルチ GPU、あるいはクラウド TPU 上に彼らの訓練をスケールすることを望むかもしれません。
これらのユースケースをサポートするために、利用可能な 4 つのストラテジーを持ちます。次のセクションでは現時点で TF 2.0-alpha 内でどのシナリオにおいてこれらのどれがサポートされるかについて話しをします。
MirroredStrategy
tf.distribute.MirroredStrategy は一つのマシン上のマルチ GPU 上で同期分散訓練をサポートします。それは GPU デバイス毎に一つのレプリカを作成します。モデルの各変数は総てのレプリカに渡りミラーリングされます。一緒に、これらの変数は MirroredVariable と呼ばれる単一の概念的な変数を形成します。これらの変数は同一の更新を適用することにより互いに同期されます。
効率的な all-reduce アルゴリズムはデバイスを通して変数更新を伝達するために使用されます。all-reduce は総てのデバイスに渡り tensor をそれらを合計することにより集めてそれらを各デバイスで利用可能にします。それは融合されたアルゴリズムで非常に効率的で同期のオーバーヘッドを著しく減らせます。デバイス間で利用可能な通信のタイプに依拠して、利用可能な多くの all-reduce アルゴリズムと実装があります。デフォルトでは、それは all-reduce 実装として NVIDIA NCCL を使用します。ユーザはまた提供する 2, 3 の他のオプション間で選択できて、彼ら自身のものを書くこともできます。
ここに MirroredStrategy を作成する最も単純な方法があります :
mirrored_strategy = tf.distribute.MirroredStrategy()
WARNING: Logging before flag parsing goes to stderr. W0307 18:06:27.818398 139925123892992 cross_device_ops.py:1111] Not all devices in `tf.distribute.Strategy` are visible to TensorFlow.
これは MirroredStrategy インスタンスを作成します、これは TensorFlow に見える総ての GPU を使用して、そして交差デバイス通信として NCCL を使用します。
貴方のマシン上の GPU の一部だけを利用することを望む場合、それをこのようにできます :
mirrored_strategy = tf.distribute.MirroredStrategy(devices=["/gpu:0", "/gpu:1"])
W0307 18:06:27.835354 139925123892992 cross_device_ops.py:1106] Not all devices in `tf.distribute.Strategy` are visible to TensorFlow.
交差デバイス通信をオーバーライドすることを望む場合、cross_device_ops 引数を使用して tf.distribute.CrossDeviceOps のインスタンスを供給することによりそれを行なうことができます。デフォルトの tf.distribute.NcclAllReduce 以外の 2 つの他のオプションとして現在 tf.distribute.HierarchicalCopyAllReduce と tf.distribute.ReductionToOneDevice を提供しています。
mirrored_strategy = tf.distribute.MirroredStrategy(
cross_device_ops=tf.distribute.HierarchicalCopyAllReduce())
W0307 18:06:27.850420 139925123892992 cross_device_ops.py:1111] Not all devices in `tf.distribute.Strategy` are visible to TensorFlow.
MultiWorkerMirroredStrategy
tf.distribute.experimental.MultiWorkerMirroredStrategy は MirroredStrategy に非常に類似しています。それはマルチ・ワーカーに渡る同期分散訓練を実装し、それぞれ潜在的にマルチ GPU を持ちます。MirroredStrategy と同様に、それはモデルの総ての変数のコピーを総てのワーカーに渡り各デバイス上に作成します。
それは CollectiveOps を変数の同期を取るために使用されるマルチワーカー all-reduce 通信メソッドとして使用します。collective op は TensorFlow グラフの単一 op で、それはハードウェア、ネットワーク・トポロジーそして tensor サイズに従って TensorFlow ランタイムで all-reduce アルゴリズムを自動的に選択できます。
それは追加のパフォーマンス最適化も実装します。例えば、それは小さい tensor 上のマルチ all-reductions をより巨大な tensor 上のより少ない all-reductions に変換する静的最適化を含みます。更に、それをプラグイン・アーキテクチャを持つように設計していますので、将来的にユーザは彼らのハードウェアのためにより良く調整されたアルゴリズムをプラグインできるでしょう。collective ops はまた broadcast と all-gather のような他の collective 演算も実装することに注意してください。
ここに MultiWorkerMirroredStrategy を作成する最も単純な方法があります :
multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
W0307 18:06:27.865478 139925123892992 cross_device_ops.py:1111] Not all devices in `tf.distribute.Strategy` are visible to TensorFlow.
MultiWorkerMirroredStrategy は現在 collective ops の 2 つの異なる実装間で選択することを可能にします。CollectiveCommunication.RING は gRPC を通信層として使用して ring-based collectives を実装します。CollectiveCommunication.NCCL は collectives を実装するために Nvidia の NCCL を使用します。CollectiveCommunication.AUTO は選択を実行時まで延期します。collective 実装の最善の選択は GPU の数と種類、そしてクラスタのネットワーク相互接続に依拠します。それらをこのように指定できます :
multiworker_strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy(
tf.distribute.experimental.CollectiveCommunication.NCCL)
W0307 18:06:27.882048 139925123892992 cross_device_ops.py:1111] Not all devices in `tf.distribute.Strategy` are visible to TensorFlow.
マルチ GPU 訓練と比較して、マルチワーカー訓練を進めるための主要な違いの一つはマルチワーカー・セットアップです。”TF_CONFIG” 環境変数はクラスタの一部である各ワーカーへのクラスタ構成を指定するための TensorFlow の標準的な方法です。どのようにこれが成されるかのより詳細については下の “TF_CONFIG” セクションを見てください。
Note: このストラテジーは 実験的 です、何故ならば現在それを改良してより多くのシナリオで動作するようにしているからです。その一環として、 将来的に API が変更されることを想定してください。
TPUStrategy
tf.distribute.experimental.TPUStrategy はユーザに彼らの TensorFlow 訓練を Tensor 処理ユニット (TPU) 上で実行させます。TPU は機械学習ワークロードを劇的に高速化するように設計された Google の特化された ASIC です。それらは Google Colab、TensorFlow Research Cloud そして Google Compute Engine で利用可能です。
分散訓練アーキテクチャの観点からは、TPUStrategy は同じ MirroredStrategy です – それは同期分散訓練を実装します。TPU はマルチ TPU コアに渡る効率的な all-reduce のそれら自身の実装と他の collective 演算を提供します、それらは TPUStrategy で使用されます。
ここにどのように TPUStrategy をインスタンス化するかがあります。Note: このコードを Colab で実行するには、Colab 実行時に TPU を選択すべきです。実行可能なバージョンについては Using TPUs ガイド (訳注: リンク切れ) を見てください。
resolver = tf.distribute.cluster_resolver.TPUClusterResolver() tf.tpu.experimental.initialize_tpu_system(resolver) tpu_strategy = tf.distribute.experimental.TPUStrategy(resolver)
TPUClusterResolver インスタンスは TPU を見つける助けをします。Colab では、それにどのような引数も指定する必要がありません。これを Cloud TPU で使用することを望む場合、tpu 引数で貴方の TPU リソースの名前を指定する必要があります。またプログラムの最初で tpu システムを明示的に初期化する必要もあります。これは TPU が計算のために使用できる前に必要で理想的には最初に成されるべきです、何故ならばそれはまた TPU メモリも拭き取るために総ての状態が失われるからです。
ParameterServerStrategy
tf.distribute.experimental.ParameterServerStrategy はパラメータサーバ訓練をサポートします。それはマルチ GPU 同期ローカル訓練か非同期マルチマシン訓練のために使用できます。一つのマシン上のローカルで訓練するために使用されるとき、変数はミラーリングされません、代わりにそれらは CPU に置かれて演算は総てのローカル GPU に渡り複製されます。マルチマシン設定では、幾つかのマシンはワーカーとしてそして幾つかはパラメータサーバとして指定されます。モデルの各変数は一つのパラメータサーバに置かれます。計算は総てのワーカーの総ての GPU に渡り複製されます。
コードの観点からは、それは他のストラテジーと同様に見えます :
ps_strategy = tf.distribute.experimental.ParameterServerStrategy()
マルチワーカー訓練については、”TF_CONFIG” はクラスタのパラメータサーバとワーカーの構成を指定する必要があります、それについて下の “TF_CONFIG” で更に読むことができます。
ここまで利用可能な異なるストラテジーが何であるかそしてそれらをどのようにインスタンス化できるかについて話してきました。次の幾つかのセクションでは、貴方の訓練を分散するためにそれらを使用する異なる方法について話します。このガイドでは短いコードスニペットを示して、end-to-end で実行できる完全なチュートリアルへリンクします。
Keras で tf.distribute.Strategy を使用する
私達は tf.distribute.Strategy を tf.keras に統合しました、これは Keras API 仕様の TensorFlow 実装です。tf.keras はモデルを構築して訓練するための高位 API です。
tf.keras バックエンドに統合することで、Keras ユーザに Keras 訓練フレームワークで書かれた彼らの訓練を分散することをシームレスにしました。ユーザのプログラムで変更が必要なものは以下だけです : (1) 適切な tf.distribute.Strategy のインスタンスを作成します。そして (2) Keras モデルの作成とコンパイルを strategy.scope の内側に移す。
ここに一つの dense 層を持つ非常に簡単な Keras モデルのためにこれを行なうコードのスニペットがあります :
mirrored_strategy = tf.distribute.MirroredStrategy() with mirrored_strategy.scope(): model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))]) model.compile(loss='mse', optimizer='sgd')
W0307 18:06:27.907308 139925123892992 cross_device_ops.py:1111] Not all devices in `tf.distribute.Strategy` are visible to TensorFlow.
この例では MirroredStrategy を使用していますのでこれをマルチ GPU を持つ (一つの) マシン上で実行できます。strategy.scope() はコードのどの部分を分散して実行するかを示しました。このスコープの内側でモデルを作成すると通常の変数の代わりにミラーリングされた変数を作成することを可能にします。このスコープ下でコンパイルするとユーザはこのストラテジーを使用してモデルを訓練することを意図していることを知ることができます。ひとたびこれがセットアップされれば、モデルを通常するように fit できます。MirroredStrategy は勾配を集める等、利用可能な GPU 上でモデルの訓練を複製する処理をします。
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(10) model.fit(dataset, epochs=2) model.evaluate(dataset)
W0307 18:06:28.070022 139925123892992 training_utils.py:1353] Expected a shuffled dataset but input dataset `x` is not shuffled. Please invoke `shuffle()` on input dataset. Epoch 1/2 10/10 [==============================] - 0s 7ms/step - loss: 2.1846 Epoch 2/2 10/10 [==============================] - 0s 2ms/step - loss: 2.0163 10/10 [==============================] - 0s 5ms/step - loss: 1.9289 1.9288513660430908
ここで訓練と評価入力を提供するために tf.data.Dataset を使用しました。numpy 配列も使用できます :
import numpy as np inputs, targets = np.ones((100, 1)), np.ones((100, 1)) model.fit(inputs, targets, epochs=2, batch_size=10)
Epoch 1/2 10/10 [==============================] - 0s 2ms/step - loss: 1.8610 Epoch 2/2 10/10 [==============================] - 0s 2ms/step - loss: 1.7176 <tensorflow.python.keras.callbacks.History at 0x7f4274ade400>
両者のケース (dataset か numpy) で、与えられた入力の各バッチはマルチレプリカの中で均等に分割されます。例えば、2 GPU で MirroredStrategy を使用する場合、size 10 の各バッチは 2 GPU の間で分割されて、それぞれが各ステップで 5 入力サンプルを受け取ります。それで貴方がより多くの GPU を追加するにつれて各エポックはより高速に訓練されます。典型的には、貴方がより多くのアクセラレータを追加するとき特別な計算パワーを効果的に利用するためにバッチサイズを増やすことを望むでしょう、またモデルに依拠して、学習率を再調整する必要もあるでしょう。レプリカの数を得るために strategy.num_replicas_in_sync を使用できます。
# Compute global batch size using number of replicas.
BATCH_SIZE_PER_REPLICA = 5
global_batch_size = (BATCH_SIZE_PER_REPLICA *
mirrored_strategy.num_replicas_in_sync)
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100)
dataset = dataset.batch(global_batch_size)
LEARNING_RATES_BY_BATCH_SIZE = {5: 0.1, 10: 0.15}
learning_rate = LEARNING_RATES_BY_BATCH_SIZE[global_batch_size]
今は何がサポートされますか？
TF 2.0 alpha リリースでは、MirroredStrategy を使用した Keras による訓練と、ParameterServerStrategy を使用した一つのマシンのパラメータサーバをサポートします。他のストラテジーのためのサポートは間もなくです。API とどのように使用するかも上と正確に同じです。TF 2.0 の Keras で TPUStrategy や MultiWorkerMirorredStrategy のような他のストラテジーを使用することを望む場合、現在は eager execution を無効にする (tf.compat.v1.disable_eager_execution()) ことによりそれを行なうことができます。注意すべきもう一つのことは Keras でマルチワーカーのために MultiWorkerMirorredStrategy を使用するとき、現在ユーザは異なるワーカーのためにデータを明示的にシャードするかシャッフルしなければならないことですが、将来的には知的に入力データを自動的にシャードするようにこれを変更します。
サンプルとチュートリアル
ここに Keras の end-to-end で上の統合を示すチュートリアルとサンプルのリストがあります : 1. MirroredStrategy で MNIST を訓練する チュートリアル。2. TPUStrategy で Fashion MNIST を訓練する チュートリアル (訳注: リンク切れ) (現在は disable_eager_execution を使用)。3. MirroredStrategy を使用して ImageNet データで公式 ResNet50 訓練。4. TPUStrategy で Cloud TPU 上 Imagenet データで訓練された ResNet50。このサンプルは現在は TensorFlow 1.x で動作のみ動作することに注意してください。
Estimator で tf.distribute.Strategy を使用する
tf.estimator は分散訓練 TensorFlow API で、それは非同期パラメータサーバのアプローチを元々サポートしていました。Keras でのように、私達は tf.distribute.Strategy を tf.Estimator に統合しました、その結果 Estimator を訓練に使用しているユーザは彼らのコードへの僅かの変更で訓練を容易に分散に変更できます。これにより、estimater ユーザは今では TPU の使用に加えて、マルチ GPU とマルチワーカー上で同期分散訓練を行なうことができます。
tf.distribute.Strategy の使用方法は Keras の場合と僅かに異なります。strategy.scope を使用する代わりに、今は strategy オブジェクトを Estimator のための RunConfig に渡します。
ここに premade estimator LinearRegressor と MirroredStrategy でこれを示すコードスニペットがあります :
mirrored_strategy = tf.distribute.MirroredStrategy()
config = tf.estimator.RunConfig(
train_distribute=mirrored_strategy, eval_distribute=mirrored_strategy)
regressor = tf.estimator.LinearRegressor(
feature_columns=[tf.feature_column.numeric_column('feats')],
optimizer='SGD',
config=config)
W0307 18:06:33.035279 139925123892992 cross_device_ops.py:1111] Not all devices in `tf.distribute.Strategy` are visible to TensorFlow. W0307 18:06:33.037824 139925123892992 estimator.py:1799] Using temporary folder as model directory: /tmp/tmp9yy5cqbb
ここでは premade Estimator を使用しますが、同じコードがカスタム Estimator でもまた動作します。train_evaluate は訓練がどのように分散されるかを決定し、そして eval_distribute は評価がどのように分散されるかを決定します。これはもう一つの Keras との違いです、そこでは訓練と評価の両者のために同じ strategy を使用します。
今ではこの Estimator を input 関数で訓練して評価することができます :
def input_fn():
dataset = tf.data.Dataset.from_tensors(({"feats":[1.]}, [1.]))
return dataset.repeat(1000).batch(10)
regressor.train(input_fn=input_fn, steps=10)
regressor.evaluate(input_fn=input_fn, steps=10)
W0307 18:06:33.136213 139916847593216 distribute_lib.py:830] Partitioned variables are disabled when using current tf.distribute.Strategy.
W0307 18:06:33.151402 139916847593216 deprecation.py:323] From /usr/local/lib/python3.5/dist-packages/tensorflow/python/feature_column/feature_column_v2.py:2758: to_float (from tensorflow.python.ops.math_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.cast` instead.
W0307 18:06:34.565543 139925123892992 deprecation.py:323] From /usr/local/lib/python3.5/dist-packages/tensorflow/python/training/saver.py:1276: checkpoint_exists (from tensorflow.python.training.checkpoint_management) is deprecated and will be removed in a future version.
Instructions for updating:
Use standard file APIs to check for files with this prefix.
{'average_loss': 1.4210855e-14,
'global_step': 10,
'label/mean': 1.0,
'loss': 1.4210855e-14,
'prediction/mean': 0.99999994}
ここで強調したい Estimator と Keras 間のもう一つの違いは入力処理です。Keras では、dataset の各バッチはマルチ・レプリカに渡り分割されることに言及しました。Estimator では、けれども、ユーザは input_fn を提供して彼らのデータがワーカーとデバイスに渡りどのように分散されることを望むかについて完全な制御を持ちます。私達はバッチの自動的分割を行ないませんし、異なるワーカーに渡りデータを自動的にシャードもしません。提供される input_fn はワーカー毎に一度呼び出され、ワーカー毎に一つの dataset を与えます。それからその dastaset から一つのバッチがそのワーカー上の一つのレプリカに供給され、従って 1 ワーカー上の N レプリカのために N バッチを消費します。換言すれば、input_fn から返される dataset はサイズ PER_REPLICA_BATCH_SIZE のバッチを提供するべきです。そしてステップのためのグローバル・バッチサイズは PER_REPLICA_BATCH_SIZE * strategy.num_replicas_in_sync として得られます。マルチワーカー訓練を行なうとき、ユーザはまたデータをワーカーに渡り分割するかそれぞれの上でランダムシードでシャッフルすることもまた望むでしょう。これをどのように行なうかのサンプルを マルチワーカー・チュートリアル で見ることができます。
Estimator で MirroredStrategy を使用するサンプルを示しました。貴方はまた Estimator で TPUStrategy を使用することもまたできます、正確に同じ方法で :
config = tf.estimator.RunConfig(
train_distribute=tpu_strategy, eval_distribute=tpu_strategy)
そして同様に、マルチワーカーとパラメータサーバのストラテジーもまた使用できます。コードは同じままですが、tf.estimator.train_and_evaluate を使用してクラスタで動作する各バイナリに対して “TF_CONFIG” 環境変数を設定する必要があります。
今は何がサポートされますか？
TF 2.0 alpha リリースでは、総てのストラテジーを使用して Estimater による訓練をサポートします。
サンプルとチュートリアル
ここに幾つかのサンプルがあります、これらは Estimator で様々なストラテジーの end-to-end の使用方法を示します :
- MultiWorkerMirroredStrategy を使用してマルチワーカーで MNIST を訓練する [Tutorial]((../tutorials/distribute/multi_worker)
- Kuberentes テンプレートを使用した tensorflow/ecosystem のマルチワーカー訓練のための end-to-end サンプル。このサンプルは Keras モデルから始めてそれを tf.keras.estimator.model_to_estimator API を使用して Estimator に変換します。
- 公式 ResNet50 モデル、これは MirroredStrategy か MultiWorkerMirroredStrategy を使用して訓練できます。
- TPUStrategy による ResNet50 サンプル。
カスタム訓練ループで tf.distribute.Strategy を使用する
貴方が見たように、高位 API で tf.distrbute.Strategy を使用するのは 2, 3 行のコード変更だけです。もう少しの努力で、tf.distrbute.Strategy はこれらのフレームワークを使用していない他のユーザによっても使用できます。
TensorFlow は様々な種類のユースケースのために使用されて (研究者のような) あるユーザは彼らの訓練ループに渡るより多くの柔軟性と制御を必要とします。これは彼らに Estimator や Keras のような高位フレームワークを使用することを困難にします。例えば、GAN を使用するある人は各ラウンドで異なる数の generator や discriminator ステップを取ることを望むかもしれません。同様に、高位フレームワークは強化学習訓練のためにはそれほど適してはいません。そのためこれらのユーザは通常は彼ら自身の訓練ループを書きます。
これらのユーザのために、tf.distrbute.Strategy クラスを通してメソッドのコアセットを提供します。これらの使用は初期化でコードの少量の再構成が必要かもしれませんが、ひとたびそれが成されれば、strategy インスタンスを単に変更することによりユーザは GPU / TPU / マルチマシンの間で切り替えることができるはずです。
ここで前と同じ Keras モデルを使用して単一の訓練サンプルのためのこのユースケースを表わす簡単なスニペットを示します。Note: これらの API は依然として実験的で TensorFlow 2.0 でそれらをよりユーザフレンドリーにするために改良しています。
最初に、strategy のスコープの内側でモデルと optimizer を作成します。これはこのモデルと optimizer で作成された任意の変数がミラーリングされる変数であることを確かなものにします。
with mirrored_strategy.scope(): model = tf.keras.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))]) optimizer = tf.keras.optimizers.SGD()
次に、入力 dataset を作成して strategy に基づいて dataset を分散するために make_dataset_iterator を呼び出します。この API は近い将来に変更されることが予定されています。
with mirrored_strategy.scope():
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(1000).batch(
global_batch_size)
input_iterator = mirrored_strategy.make_dataset_iterator(dataset)
それから、訓練のワンステップを定義します。勾配を計算するために tf.GradientTape を使用してモデルの変数を更新するためにそれらの勾配を適用するために optimizer を使用します。この訓練ループを分散するために、関数 step_fn を中に入れてそれを前に作成された iterator と一緒に strategy.experimental_run に渡します :
@tf.function
def train_step():
def step_fn(inputs):
features, labels = inputs
with tf.GradientTape() as tape:
logits = model(features)
cross_entropy = tf.nn.softmax_cross_entropy_with_logits(
logits=logits, labels=labels)
loss = tf.reduce_sum(cross_entropy) * (1.0 / global_batch_size)
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(list(zip(grads, model.trainable_variables)))
return loss
per_replica_losses = mirrored_strategy.experimental_run(
step_fn, input_iterator)
mean_loss = mirrored_strategy.reduce(
tf.distribute.ReduceOp.MEAN, per_replica_losses)
return mean_loss
上のコードで注意すべき 2, 3 の他のこと: 1. 損失を計算するために tf.nn.softmax_cross_entropy_with_logits を使用しました。そしてそれからトータル損失をグローバル・バッチサイズでスケールしました。これは重要です、何故ならば総てのレプリカは同期して訓練して訓練の各ステップのサンプル数はグローバル・バッチだからです。もし tf.losses や tf.keras.losses から TensorFlow の標準損失を使用していれば、それらは分散 aware で strategy がスコープにあるときはいつでもレプリカ数によるスケーリングを処理します。2. experimental_run により返される結果を集計するために strategy.reduce API を使用しました。experimental_run は strategy の各ローカル・レプリカからの結果を返して、そしてこれらの結果を消費するための複数の方法があります。集計値を得るためにそれらを reduce できます。ローカル・レプリカ毎に一つ、結果に含まれる値のリストを得るために strategy.unwrap(results)* を行なうこともできます。
* 変更される予定です。
最後に、ひとたび訓練ステップを定義したならば、 iterator を初期化してループで訓練を実行できます :
with mirrored_strategy.scope():
input_iterator.initialize()
for _ in range(10):
print(train_step())
tf.Tensor(0.0, shape=(), dtype=float32) tf.Tensor(0.0, shape=(), dtype=float32) tf.Tensor(0.0, shape=(), dtype=float32) tf.Tensor(0.0, shape=(), dtype=float32) tf.Tensor(0.0, shape=(), dtype=float32) tf.Tensor(0.0, shape=(), dtype=float32) tf.Tensor(0.0, shape=(), dtype=float32) tf.Tensor(0.0, shape=(), dtype=float32) tf.Tensor(0.0, shape=(), dtype=float32) tf.Tensor(0.0, shape=(), dtype=float32)
上の例では、入力を貴方の訓練に提供するために make_dataset_iterator を使用しました。私達はまた 2 つの追加の API を提供します : 他の種類の入力をサポートするための make_input_fn_iterator と make_experimental_numpy_iterator。tf.distribute.Strategy のそれらのドキュメントを見てくださいそしてそれらが make_dataset_iterator とどのように違うのかを。
これはカスタム訓練ループを分散するために tf.distribute.Strategy API を使用する最も単純なケースをカバーします。私達はこれらの API を改良する過程にあります。このユースケースはユーザ側のより多くのワークを必要としますので、将来的にこのユースケースのための個別の詳細なガイドを公開します。
今は何がサポートされますか？
TF 2.0 alpha リリースでは、上で示されたように MirroredStrategy を使用してカスタム訓練ループを持つ訓練をサポートします。他のストラテジーのサポートは間もなくです。カスタム訓練ループで TF 2.0 で TPUStrategy のような他のストラテジーを使用することを望むのであれば、現在 eager execution を無効にすることによりそれを行なうことができます (tf.compat.v1.disable_eager_execution())。コードは同様のままです、訓練を実行するために TF 1.x グラフとセッションを使用する必要があることを除いて。MultiWorkerMirorredStrategy サポートは将来的にもたらされます。
サンプルとチュートリアル
ここにカスタム訓練ループで分散ストラテジーを使用するためのサンプルがあります : 1. MirroredStrategy を使用して MNIST を訓練する チュートリアル。2. MirroredStrategy を使用する DenseNet サンプル。
他のトピック
このセクションでは、複数のユースケースに関連する幾つかのトピックをカバーします。
TF_CONFIG 環境変数をセットアップする
マルチワーカー訓練のためには、前に言及したように、貴方のクラスタの各バイナリ実行のために “TF_CONFIG” 環境変数を設定する必要があります。”TF_CONFIG” 環境変数は JSON 文字列でこれは何のタスクがクラスタを構成するか、それらのアドレスとクラスタの各タスクの役割りを指定します。tensorflow/ecosystem レポで Kubernetes テンプレートを提供します、これは貴方の訓練タスクのために “TF_CONFIG” を設定します。
“TF_CONFIG” の一つの例は :
os.environ["TF_CONFIG"] = json.dumps({
"cluster": {
"worker": ["host1:port", "host2:port", "host3:port"],
"ps": ["host4:port", "host5:port"]
},
"task": {"type": "worker", "index": 1}
})
この “TF_CONFIG” はクラスタに 3 つのワーカーと 2 つの ps タスクがそれらのホストとポートと一緒にあることを指定しています。”task” パートはクラスタの現在のタスクの役割り, ワーカー 1 (2 番目のワーカー) を指定しています。クラスタの有効な役割りは “chief”, “worker”, “ps” と “evaluator” です。tf.distribute.experimental.ParameterServerStrategy を使用するとき以外は “ps” ジョブはないはずです。
以上