Book a Demo!
CoCalc Logo Icon
StoreFeaturesDocsShareSupportNewsAboutPoliciesSign UpSign In
tensorflow
GitHub Repository: tensorflow/docs-l10n
 Path: blob/master/site/ja/guide/migrate/tpu_estimator.ipynb
25118 views
Kernel: Python 3
#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

このガイドでは、TPU で実行されているワークフローを TensorFlow 1 の TPUEstimator API から TensorFlow 2 の TPUStrategy API に移行する方法を示します。

  • TensorFlow 1 では、tf.compat.v1.estimator.tpu.TPUEstimator API を使用して、モデルのトレーニングと評価、推論の実行、モデルの（サービングのための）保存を（クラウド）TPU で行うことができます。

  • TensorFlow 2 で、TPU と TPU Pod（専用の高速ネットワーク インターフェースで接続された TPU デバイスのコレクション）で同期トレーニングを実行するには、TPU 分散ストラテジー（tf.distribute.TPUStrategy）を使用する必要があります。このストラテジーは、モデル構築（tf.keras.Model）、オプティマイザ（tf.keras.optimizers.Optimizer）、およびトレーニング（Model.fit）、およびカスタムトレーニングループ（tf.function と tf.GradientTape を使用）で使用できます。

エンドツーエンドの TensorFlow 2 の例については、TPU の使用ガイド（TPU の分類セクション）と TPU で BERT を使用して GLUE タスクを解決するチュートリアルを参照してください。また、分散トレーニングガイドも役立つかもしれません。このガイドでは、TPUStrategy を含むすべての TensorFlow 分散ストラテジーについて説明されています。

セットアップ

まず、インポートし、デモ用の単純なデータセットから始めます。

import tensorflow as tf
import tensorflow.compat.v1 as tf1

features = [[1., 1.5]]
labels = [[0.3]]
eval_features = [[4., 4.5]]
eval_labels = [[0.8]]

TensorFlow 1: TPUEstimator を使用して TPU でモデルを駆動する

このセクションでは、TensorFlow 1 で tf.compat.v1.estimator.tpu.TPUEstimator を使用してトレーニングと評価を実行する方法を示します。

TPUEstimator を使用するには、まず、トレーニングデータの入力関数、評価データの評価入力関数、および特徴量とラベルを使用してトレーニング演算がどのように定義されるかを Estimator に伝えるモデル関数など、いくつかの関数を定義します。

def _input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((features, labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _eval_input_fn(params):
  dataset = tf1.data.Dataset.from_tensor_slices((eval_features, eval_labels))
  dataset = dataset.repeat()
  return dataset.batch(params['batch_size'], drop_remainder=True)

def _model_fn(features, labels, mode, params):
  logits = tf1.layers.Dense(1)(features)
  loss = tf1.losses.mean_squared_error(labels=labels, predictions=logits)
  optimizer = tf1.train.AdagradOptimizer(0.05)
  train_op = optimizer.minimize(loss, global_step=tf1.train.get_global_step())
  return tf1.estimator.tpu.TPUEstimatorSpec(mode, loss=loss, train_op=train_op)

これらの関数を定義したら、クラスタ情報を提供する tf.distribute.cluster_resolver.TPUClusterResolver と tf.compat.v1.estimator.tpu.RunConfig オブジェクトを作成します。定義したモデル関数に加えて、TPUEstimator を作成できるようになりました。ここでは、チェックポイントの保存をスキップしてフローを簡素化します。次に、TPUEstimator のトレーニングと評価の両方のバッチサイズを指定します。

cluster_resolver = tf1.distribute.cluster_resolver.TPUClusterResolver(tpu='')
print("All devices: ", tf1.config.list_logical_devices('TPU'))

tpu_config = tf1.estimator.tpu.TPUConfig(iterations_per_loop=10)
config = tf1.estimator.tpu.RunConfig(
    cluster=cluster_resolver,
    save_checkpoints_steps=None,
    tpu_config=tpu_config)
estimator = tf1.estimator.tpu.TPUEstimator(
    model_fn=_model_fn,
    config=config,
    train_batch_size=8,
    eval_batch_size=8)

TPUEstimator.train を呼び出して、モデルのトレーニングを開始します。

estimator.train(_input_fn, steps=1)

次に、TPUEstimator.evaluate を呼び出して、評価データを使用してモデルを評価します。

estimator.evaluate(_eval_input_fn, steps=1)

TensorFlow 2: Keras Model.fit と TPUStrategy を使用して TPU でモデルを駆動する

TensorFlow 2 で TPU ワーカーをトレーニングするには、tf.distribute.TPUStrategy を Keras API とともに使用して、モデルの定義とトレーニング/評価を行います。（Keras Model.fit とカスタムトレーニングループ（tf.function と tf.GradientTape を使用）を使用したトレーニングのその他の例については、TPU の使用ガイドを参照してください。）

リモートクラスタに接続して TPU ワーカーを初期化するには初期化作業を行う必要があります。まず TPUClusterResolver を作成してクラスタ情報を提供し、クラスタに接続します。（詳しくは、TPU の使用ガイドの TPU 初期化 セクションを参照してください）。

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
print("All devices: ", tf.config.list_logical_devices('TPU'))

次に、データの準備が整ったら、TPUStrategy を作成し、このストラテジーの範囲でモデル、指標、オプティマイザを定義します。

TPUStrategy で同等のトレーニング速度を達成するには、各 tf.function 呼び出し中に実行されるバッチ数を指定するために Model.compile で steps_per_execution の数値を選択する必要があります。これは、パフォーマンスにとって重要です。この引数は、TPUEstimator で使用される iterations_per_loop に似ています。カスタムトレーニングループを使用している場合は、tf.function を使用したトレーニング関数内で複数のステップが実行されていることを確認する必要があります。詳細については、TPU の使用ガイドの tf.function 内の複数のステップによるパフォーマンスの改善セクションに移動してください。

tf.distribute.TPUStrategy は、動的形状計算の上限を推測できる境界のある動的形状をサポートします。ただし、動的形状では、静的形状と比較してパフォーマンスのオーバーヘッドが発生する場合があります。そのため、特にトレーニングでは、可能であれば入力形状を静的にすることを推奨します。ストリームに残っているサンプルの数がバッチ サイズよりも少ない可能性があるため、動的な形状を返す 1 つの一般的な演算は tf.data.Dataset.batch(batch_size) です。したがって、TPU でトレーニングする場合は、最高のトレーニングのパフォーマンスを得るために tf.data.Dataset.batch(..., drop_reminder=True) を使用する必要があります。

dataset = tf.data.Dataset.from_tensor_slices(
    (features, labels)).shuffle(10).repeat().batch(
        8, drop_remainder=True).prefetch(2)
eval_dataset = tf.data.Dataset.from_tensor_slices(
    (eval_features, eval_labels)).batch(1, drop_remainder=True)

strategy = tf.distribute.TPUStrategy(cluster_resolver)
with strategy.scope():
  model = tf.keras.models.Sequential([tf.keras.layers.Dense(1)])
  optimizer = tf.keras.optimizers.Adagrad(learning_rate=0.05)
  model.compile(optimizer, "mse", steps_per_execution=10)

以上でトレーニングデータセットを使用してモデルをトレーニングする準備が整いました。

model.fit(dataset, epochs=5, steps_per_epoch=10)

最後に、評価データセットを使用してモデルを評価します。

model.evaluate(eval_dataset, return_dict=True)

次のステップ

TensorFlow 2 の TPUStrategy についての詳細は、次のリソースを参照してください。

  • ガイド: TPU の使用（Keras Model.fit を使用したトレーニング / tf.distribute.TPUStrategy を使用したカスタムトレーニングループと、tf.function でパフォーマンスを改善するためのヒントを含む）

  • ガイド: TensorFlow による分散トレーニング

トレーニングのカスタマイズの詳細については、次を参照してください。

TPU（Google の機械学習専用 ASIC）は、Google Colab、TPU Research Cloud、Cloud TPU から入手できます。