input.ipynb

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

#@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.

分散入力

TensorFlow.org で表示

Google Colabで実行

GitHubでソースを表示

ノートブックをダウンロード

tf.distribute API は、1 台のコンピュータから複数のコンピュータにトレーニングを簡単にスケーリングする方法を提供します。モデルをスケーリングする際には、ユーザーは入力を複数のデバイスに分散する必要がありますが、tf.distribute は、入力を自動的にデバイスに分散できる API を提供します。

このガイドは、tf.distribute API を使用して、分散データセットとイテレータを作成するためのさまざまな方法を見ていきます。さらに、次のトピックについても説明しています。

• tf.distribute.Strategy.experimental_distribute_dataset と tf.distribute.Strategy.distribute_datasets_from_function の使用方法、およびこれらを使用したシャーディングとバッチオプション

• 分散データセットのさまざまなイテレーション方法

• tf.distribute.Strategy.experimental_distribute_dataset/tf.distribute.Strategy.distribute_datasets_from_function API と tf.data API の違い、および使用時の制限

このガイドでは、Keras API を使用した分散入力の使用方法は説明されていません。

分散データセット

tf.distribute API を使用してスケーリングするには、tf.data.Dataset を使って入力を表現します。tf.distribute は、パフォーマンス最適化を定期的に実装に統合しながら、tf.data.Dataset と効率的に動作します (各アクセラレータデバイスへのデータの自動プリフェッチ機能、定期的なパフォーマンスの更新など)。tf.data.Dataset 以外を使用するユースケースがある場合は、このガイドの Tensor 入力セクションを参照してください。非分散型トレーニングループでは、tf.data.Dataset インスタンスを作成してから要素をイテレートします。次に例を示します。

import tensorflow as tf

# Helper libraries
import numpy as np
import os

print(tf.__version__)

# Simulate multiple CPUs with virtual devices
N_VIRTUAL_DEVICES = 2
physical_devices = tf.config.list_physical_devices("CPU")
tf.config.set_logical_device_configuration(
    physical_devices[0], [tf.config.LogicalDeviceConfiguration() for _ in range(N_VIRTUAL_DEVICES)])

print("Available devices:")
for i, device in enumerate(tf.config.list_logical_devices()):
  print("%d) %s" % (i, device))

global_batch_size = 16
# Create a tf.data.Dataset object.
dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(global_batch_size)

@tf.function
def train_step(inputs):
  features, labels = inputs
  return labels - 0.3 * features

# Iterate over the dataset using the for..in construct.
for inputs in dataset:
  print(train_step(inputs))

ユーザーの既存のコードへの変更を最小限に抑えて tf.distribute ストラテジーを使用できるように、tf.data.Dataset インスタンスを分散し、分散化されたデータセットインスタンスを返す、2 つの API が導入されています。その分散化されたデータセットインスタンスを以前と同様にイテレートして、モデルをトレーニングすることができます。では、これら 2 つの API を詳しく見てみましょう。tf.distribute.Strategy.experimental_distribute_dataset API と tf.distribute.Strategy.distribute_datasets_from_function API です。

tf.distribute.Strategy.experimental_distribute_dataset

使い方

この API はtf.data.Datasetインスタンスを入力として取り、tf.distribute.DistributedDatasetインスタンスを返します。この入力データセットを、グローバルバッチサイズと同じ値でバッチ化します。このグローバルバッチサイズは、1 つのステップで処理する全デバイスのサンプル数です。この分散データセットのイテレーションを Python 式に行うか、iter を使用してイテレータを作成します。返されるオブジェクトはtf.data.Datasetインスタンスではなく、またデータセットを変換したり検査したりするほかの API をまったくサポートしていません。これは、入力をさまざまなレプリカにシャーディングするための特定の方法がない場合に推奨される API です。

global_batch_size = 16
mirrored_strategy = tf.distribute.MirroredStrategy()

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(global_batch_size)
# Distribute input using the `experimental_distribute_dataset`.
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)
# 1 global batch of data fed to the model in 1 step.
print(next(iter(dist_dataset)))

プロパティ

バッチ処理

tf.distribute は、グローバルバッチサイズを同期中のレプリカの数で除算した値に等しい新しいバッチサイズで、入力 tf.data.Dataset インスタンスのバッチを再作成します。同期中のレプリカの数は、トレーニング中の勾配の allreduce に参加しているデバイスの数と同等です。ユーザーが分散イテレータで next を呼び出すと、レプリカ当たりのデータサイズが各レプリカに返されます。再バッチされたデータセットのカーディナリティは、必ずレプリカ数の倍数になります。次にいくつかの例を示します。

• tf.data.Dataset.range(6).batch(4, drop_remainder=False)

  • 分散無し:

    • バッチ 1: [0, 1, 2, 3]

    • バッチ 2: [4, 5]

  • 2 つのレプリカで分散。最後のバッチ（[4, 5]）は、2つのレプリカ間で分割。

  • バッチ 1:

    • レプリカ 1: [0, 1]

    • レプリカ 2: [2, 3]

  • バッチ 2:

    • レプリカ 1: [4]

    • レプリカ 2: [5]

• tf.data.Dataset.range(4).batch(4)

  • 分散無し:

    • バッチ 1: [0, 1, 2, 3]

  • 5 つのレプリカで分散:

    • バッチ 1:

      • レプリカ 1: [0]

      • レプリカ 2: [1]

      • レプリカ 3: [2]

      • レプリカ 4: [3]

      • レプリカ 5: []

• tf.data.Dataset.range(8).batch(4)

  • 分散無し:

    • バッチ 1: [0, 1, 2, 3]

    • バッチ 2: [4, 5, 6, 7]

  • 3 つのレプリカで分散:

    • バッチ 1:

      • レプリカ 1: [0, 1]

      • レプリカ 2: [2, 3]

      • レプリカ 3: []

    • バッチ 2:

      • レプリカ 1: [4, 5]

      • レプリカ 2: [6, 7]

      • レプリカ 3: []

注意: 上記の例は、異なるレプリカでグローバルバッチがどのように分割されるかのみを説明しています。実装によって実際の値が異なる可能性があるため、各レプリカで最終的に得られる可能性のある実際の値に依存することはお勧めできません。

データセットのバッチの再作成には、レプリカの数とともに直線的に増加する空間的コストがあります。つまり、マルチワーカートレーニングのユースケースで言えば、入力パイプラインで OOM エラーが発生する可能性があります。

シャーディング

tf.distribute は、MultiWorkerMirroredStrategy と TPUStrategy のマルチワーカートレーニングで入力データセットの自動シャーディングも行います。各データセットはワーカーの CPU デバイス上に作成されます。データセットを一連のワーカーで自動シャーディングすると、各ワーカーにデータセット全体のサブセットが割り当てられることになります（適切な tf.data.experimental.AutoShardPolicy が設定されている場合）。これは、各ステップにおいて、オーバーラップしていないデータセット要素のグローバルバッチサイズが各ワーカーで処理されるようにするためです。自動シャーディングには、tf.data.experimental.DistributeOptions で指定できる 2 つのオプションがあります。ParameterServerStrategy のマルチワーカーでは自動シャーディングは行われません。このストラテジーでのデータセット作成の詳細については、ParameterServerStrategy のチュートリアルをご覧ください。

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(64).batch(16)
options = tf.data.Options()
options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
dataset = dataset.with_options(options)

tf.data.experimental.AutoShardPolicy に設定できるオプションには 3 つあります。

• AUTO: デフォルトのオプションです。ファイルごとにシャーディングしようとします。ファイルベースのデータセットが検出されない場合、ファイルごとのシャーディングは失敗し、tf.distribute はデータごとのシャーディングに切り替えます。入力データセットがファイルベースであっても、ファイル数がワーカー数より少ない場合は、InvalidArgumentError が発生します。これが発生した場合は、ポリシーを明示的にAutoShardPolicy.DATA に設定するか、入力ソースを小さなファイルに分割して、ファイル数がワーカー数よりも多くなるようにします。

• FILE: すべてのワーカーで入力をシャーディングする場合のオプションです。入力ファイルの数がワーカー数を大きく上回り、ファイル内のデータが均等に分散されている場合は、このオプションを使用します。このオプションの欠点は、ファイル内のデータが均等に分散されていない場合にアイドル状態のワーカーが存在することにあります。ファイル数がワーカー数より少ない場合、InvalidArgumentError が発生します。 これが発生した場合は、ポリシーを明示的に AutoShardPolicy.DATA に設定してください。例として、2 つのファイルをそれぞれに 1 つのレプリカを持つ 2 つのワーカーに分散します。ファイル 1 には [0, 1, 2, 3, 4, 5]、ファイル 2 には [6, 7, 8, 9, 10, 11] が含まれます。同期中のレプリカの合計数を 2 、グローバルバッチサイズを 4 とします。

  • ワーカー 0:

    • バッチ 1 = レプリカ 1: [0, 1]

    • バッチ 2 = レプリカ 1: [2, 3]

    • バッチ 3 = レプリカ 1: [4]

    • バッチ 4 = レプリカ 1: [5]

  • ワーカー 1:

    • バッチ 1 = レプリカ 2: [6, 7]

    • バッチ 2 = レプリカ 2: [8, 9]

    • バッチ 3 = レプリカ 2: [10]

    • バッチ 4 = レプリカ 2: [11]

• DATA: すべてのワーカーで要素を自動シャーディングします。各ワーカーはデータセット全体を読み取って、それに割り当てられたシャードのみを処理し、その他すべてのシャードは破棄されます。これは通常、入力ファイルの数がワーカー数より少なく、すべてのワーカー間でデータのシャーディングをより最適に行う場合に使用されます。欠点は、各ワーカーでデータセット全体が読み取られることです。例として、1 つのファイルを 2 つのワーカーで分散します。ファイル 1 には [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] が含まれます。同期中のレプリカの合計数を 2 とします。

  • ワーカー 0:

    • バッチ 1 = レプリカ 1: [0, 1]

    • バッチ 2 = レプリカ 1: [4, 5]

    • バッチ 3 = レプリカ 1: [8, 9]

  • ワーカー 1:

    • バッチ 1 = レプリカ 2: [2, 3]

    • バッチ 2 = レプリカ 2: [6, 7]

    • バッチ 3 = レプリカ 2: [10, 11]

• OFF: 自動シャーディングをオフにすると、各ワーカーはすべてのデータを処理します。例として、1 つのファイルを 2 つのワーカーで分散します。ファイル 1 には、[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11] が含まれます。同期中のレプリカの合計数を 2 とします。各ワーカーは、次のような分散になります。

  • ワーカー 0:

    • バッチ 1 = レプリカ 1: [0, 1]

    • バッチ 2 = レプリカ 1: [2, 3]

    • バッチ 3 = レプリカ 1: [4, 5]

    • バッチ 4 = レプリカ 1: [6, 7]

    • バッチ 5 = レプリカ 1: [8, 9]

    • バッチ 6 = レプリカ 1: [10, 11]

  • ワーカー 1:

    • バッチ 1 = レプリカ 2: [0, 1]

    • バッチ 2 = レプリカ 2: [2, 3]

    • バッチ 3 = レプリカ 2: [4, 5]

    • バッチ 4 = レプリカ 2: [6, 7]

    • バッチ 5 = レプリカ 2: [8, 9]

    • バッチ 6 = レプリカ 2: [10, 11]

プリフェッチ

デフォルトでは、tf.distributeはユーザーが提供するtf.data.Datasetインスタンスにプリフェッチ変換を追加します。プリフェッチ変換に対する引数buffer_sizeは同期中のレプリカの数と同等です。

tf.distribute.Strategy.distribute_datasets_from_function

使い方

この API は、入力関数を取って tf.distribute.DistributedDataset インスタンスを返します。ユーザーが渡す入力関数には tf.distribute.InputContext 引数があり、tf.data.Dataset インスタンスを返します。この API を使用すると、tf.distribute は、入力関数から返されたユーザーの tf.data.Dataset インスタンスにそれ以降の変更を適用しません。そのため、ユーザーがデータセットをバッチ処理してシャーディングする必要があります。tf.distribute は各ワーカーの CPU デバイスで入力関数を呼び出します。ユーザーが独自のバッチングとシャーディングのロジックを指定できるほか、この API は、マルチワーカートレーニングに使用される場合に、tf.distribute.Strategy.experimental_distribute_dataset よりも優れたスケーラビリティとパフォーマンスを示します。

mirrored_strategy = tf.distribute.MirroredStrategy()

def dataset_fn(input_context):
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(64).batch(16)
  dataset = dataset.shard(
      input_context.num_input_pipelines, input_context.input_pipeline_id)
  dataset = dataset.batch(batch_size)
  dataset = dataset.prefetch(2)  # This prefetches 2 batches per device.
  return dataset

dist_dataset = mirrored_strategy.distribute_datasets_from_function(dataset_fn)

プロパティ

バッチ処理

入力関数の戻り値であるtf.data.Datasetインスタンスは、レプリカごとのバッチサイズを使用してバッチ処理する必要があります。レプリカごとのバッチサイズは、グローバルバッチサイズを同期型トレーニングに参加しているレプリカの数で除算した値です。これは、tf.distributeが各ワーカーの CPU デバイスで入力関数を呼び出すためです。あるワーカーで作成されるデータセットは、そのワーカーのすべてのレプリカで使用する準備を整えています。

シャーディング

ユーザーの入力関数への引数として暗黙的に渡されるtf.distribute.InputContextオブジェクトは、内部的にtf.distributeよって作成されます。このオブジェクトには、ワーカー数、現在のワーカー ID などの情報が含まれます。この入力関数は、tf.distribute.InputContextオブジェクトの一部であるプロパティを使用し、ユーザーが設定したポリシーに従って、シャーディングを処理することができます。

プリフェッチ

tf.distribute は、ユーザーが提供する入力関数によって返される tf.data.Dataset の最後に、プリフェッチ変換を追加しません。したがって、上記の例では明示的に Dataset.prefetch を呼び出します。

注意:tf.distribute.Strategy.experimental_distribute_dataset と tf.distribute.Strategy.distribute_datasets_from_function は両方とも、tf.data.Dataset 型ではない tf.distribute.DistributedDataset インスタンスを返します。これらのインスタンスをイテレートし（「分散イテレータ」を参照）、element_spec プロパティを使用することができます。

分散イテレータ

非分散型tf.data.Datasetインスタンスと同様に、tf.distribute.DistributedDatasetインスタンスを作成してイテレートし、tf.distribute.DistributedDatasetの要素にアクセスする必要があります。次に、tf.distribute.DistributedIteratorを作成して、それをモデルのトレーニングに使用する方法を示します。

使用方法

Python 式の for ループコンストラクトを使用する

ユーザーフレンドリーな Python 式のループを使用して、tf.distribute.DistributedDataset をイテレートすることができます。tf.distribute.DistributedIterator から返される要素は、単一のtf.Tensor か、レプリカあたりの値を含む tf.distribute.DistributedValues です。tf.function にループを配置すると、パフォーマンスは上昇しますが、break とreturn は、現在、tf.function 内に配置された tf.distribute.DistributedDataset のループではサポートされていません。

global_batch_size = 16
mirrored_strategy = tf.distribute.MirroredStrategy()

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

@tf.function
def train_step(inputs):
  features, labels = inputs
  return labels - 0.3 * features

for x in dist_dataset:
  # train_step trains the model using the dataset elements
  loss = mirrored_strategy.run(train_step, args=(x,))
  print("Loss is ", loss)

iterを使用して明示的なイテレータを作成する

tf.distribute.DistributedDataset インスタンスの要素をイテレートするには、iter API を使ってtf.distribute.DistributedIterator を作成することができます。明示的なイテレータを使用すると、一定のステップ数、イテレートすることができます。tf.distribute.DistributedIterator インスタンスの dist_iterator から次の要素を取得するには、next(dist_iterator)、dist_iterator.get_next()、または dist_iterator.get_next_as_optional() を呼び出すことができます。最初の 2 つは基本的に同じです。

num_epochs = 10
steps_per_epoch = 5
for epoch in range(num_epochs):
  dist_iterator = iter(dist_dataset)
  for step in range(steps_per_epoch):
    # train_step trains the model using the dataset elements
    loss = mirrored_strategy.run(train_step, args=(next(dist_iterator),))
    # which is the same as
    # loss = mirrored_strategy.run(train_step, args=(dist_iterator.get_next(),))
    print("Loss is ", loss)

next または tf.distribute.DistributedIterator.get_next を使用すると、tf.distribute.DistributedIterator が最後に到達した場合に、OutOfRange エラーが発生します。クライアントは Python 側でそのエラーをキャッチし、チェックポイント作成や評価といった他の作業を継続することができます。ただし、次に示すようなホストトレーニングループ (tf.function ごとに複数のステップを実行する) を使用している場合は機能しません。

@tf.function
def train_fn(iterator):
  for _ in tf.range(steps_per_loop):
    strategy.run(step_fn, args=(next(iterator),))

train_fn には、tf.range 内にステップ本文をラッピングすることで、複数のステップが含まれています。この場合、ループでの依存関係のない別のイテレーションが並行して開始する可能性があるため、前のイテレーションの計算が終了する前の後の方のイテレーションで、OutOfRange エラーが発生することがあります。OutOfRange エラーが発生してしまえば、関数内のすべての演算は即座に終了されてしまいます。この状況を避ける場合は、OutOfRange エラーが発生しない別の方法として、tf.distribute.DistributedIterator.get_next_as_optional が挙げられます。get_next_as_optional は、次の要素を含む tf.experimental.Optional を返すか、tf.distribute.DistributedIterator が最後に達している場合は何の値も返しません。

# You can break the loop with `get_next_as_optional` by checking if the `Optional` contains a value
global_batch_size = 4
steps_per_loop = 5
strategy = tf.distribute.MirroredStrategy()

dataset = tf.data.Dataset.range(9).batch(global_batch_size)
distributed_iterator = iter(strategy.experimental_distribute_dataset(dataset))

@tf.function
def train_fn(distributed_iterator):
  for _ in tf.range(steps_per_loop):
    optional_data = distributed_iterator.get_next_as_optional()
    if not optional_data.has_value():
      break
    per_replica_results = strategy.run(lambda x: x, args=(optional_data.get_value(),))
    tf.print(strategy.experimental_local_results(per_replica_results))
train_fn(distributed_iterator)

element_spec プロパティを使用する

分散データセットの要素を tf.function に渡し、tf.TypeSpec の保証を必要としている場合は、tf.function の input_signature 引数を指定することができます。分散データセットの出力は、単一のデバイスまたは複数のデバイスへの入力を表せる tf.distribute.DistributedValues です。この分散値に対応する tf.TypeSpec を取得するには、tf.distribute.DistributedDataset.element_spec または tf.distribute.DistributedIterator.element_spec を使用することができます。

global_batch_size = 16
epochs = 5
steps_per_epoch = 5
mirrored_strategy = tf.distribute.MirroredStrategy()

dataset = tf.data.Dataset.from_tensors(([1.], [1.])).repeat(100).batch(global_batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

@tf.function(input_signature=[dist_dataset.element_spec])
def train_step(per_replica_inputs):
  def step_fn(inputs):
    return 2 * inputs

  return mirrored_strategy.run(step_fn, args=(per_replica_inputs,))

for _ in range(epochs):
  iterator = iter(dist_dataset)
  for _ in range(steps_per_epoch):
    output = train_step(next(iterator))
    tf.print(output)

データの前処理をする

以上、tf.data.Dataset を分散する方法を見てきましたが、データをモデルに使用する前に、データのクレンジング、変換、拡張などの前処理を行う必要があります。 次の 2 つの便利なツールを利用できます。

• Keras 前処理レイヤー: 開発者が Keras ネイティブの入力処理パイプラインを構築できるようにする一連の Keras レイヤーです。一部の Keras 前処理レイヤーには、初期化または adapt 時に設定できるトレーニング不可能な状態が含まれています (Keras 前処理レイヤー ガイドの adapt セクションを参照してください)。 ステートフルな前処理レイヤーを分散する場合、状態をすべてのワーカーに複製する必要があります。これらのレイヤーを使用するには、それらをモデルの一部にするか、データセットに適用します。

• TensorFlow Transform (tf.Transform): データ前処理パイプラインを介してインスタンスレベルとフルパスの両方のデータ変換を定義するための TensorFlow のライブラリです。Tensorflow Transform には 2 つのフェーズがあります。 1 つ目は分析フェーズです。ここでは、生のトレーニングデータがフルパスプロセスで分析され、変換に必要な統計が計算され、変換ロジックがインスタンスレベルの演算として生成されます。2 つ目は変換フェーズで、生のトレーニングデータがインスタンスレベルのプロセスで変換されます。

Keras 前処理レイヤーと Tensorflow Transform の比較

Tensorflow Transform と Keras の前処理レイヤーはどちらも、トレーニング時の前処理を分割し、推論中に前処理をモデルにバンドルして、トレーニング/サーブ スキューを減らす方法を提供します。

TFX と密に統合された Tensorflow Transform は、トレーニングパイプラインとは別のジョブで、あらゆるサイズのデータセットを分析および変換するスケーラブルな map-reduce ソリューションを提供します。単一のマシンに収まらないデータセットで分析を実行する必要がある場合は、Tensorflow Transform が最初の選択肢になります。

Keras 前処理レイヤーは、ディスクからデータを読み取った後、トレーニング時に前処理を適用する場合に最適です。これらは、Keras ライブラリのモデル開発にシームレスに適合し、adapt による小規模なデータセットの分析をサポートします。Keras 前処理レイヤーは、画像データの拡張などのユースケースをサポートし、入力データセットを通過するたびに、トレーニング用のさまざまな例が生成されます。

2 つのライブラリを混在させることもできます。この場合、Tensorflow Transform は入力データの分析と静的変換に、Keras 前処理レイヤーはトレーニング時の変換（One-Hot エンコーディングやデータ拡張など）に使用されます。

tf.distribute のベスト プラクティス

両方のツールを使用する場合、変換ロジックを初期化してデータに適用する必要があります。これにより Tensorflow リソースが作成される場合があります。これらのリソースまたは状態は、ワーカー間またはワーカーとコーディネーター間の通信を節約するために、すべてのワーカーに複製する必要があります。そのためには、他の Keras レイヤーと同じように Keras 前処理レイヤー、tft.TFTransformOutput.transform_features_layer、または tft.TransformFeaturesLayer を tf.distribute.Strategy.scope の下に作成することをお勧めします。

次の例は、tf.distribute.Strategy API を高レベルの Keras Model.fit API およびカスタムトレーニングループと別に使用する方法を示しています。

Keras 前処理レイヤーのユーザー向けの注意事項:

前処理レイヤーと大きな語彙

マルチワーカー設定で大きな語彙 (1 ギガバイト以上) を扱う場合 (tf.distribute.MultiWorkerMirroredStrategy、tf.distribute.experimental.ParameterServerStrategy、tf.distribute.TPUStrategy など)、すべてのワーカーからアクセス可能な静的ファイル (Cloud Storage などを使用する) に語彙を保存することをお勧めします。これにより、トレーニング時に語彙をすべてのワーカーに複製するのにかかる時間を短縮できます。

tf.data パイプラインでの前処理とモデルでの前処理の比較

Keras の前処理レイヤーはモデルの一部として適用することも、tf.data.Dataset に直接適用することもできますが、それぞれ利点があります。

• モデル内に前処理レイヤーを適用すると、モデルが移植可能になり、トレーニング/サービング スキューを低減するのに役立ちます。(詳細については、前処理レイヤーの使用ガイドの推論時にモデル内で前処理を行う利点セクションを参照してください)。

• tf.data パイプライン内で適用すると、プリフェッチまたは CPU へのオフロードが可能になり、アクセラレータを使用する際のパフォーマンスが向上します。

1 つ以上の TPU で実行する場合、ほとんどの場合、ユーザーは Keras 前処理レイヤーを tf.data パイプラインに配置する必要があります。すべてのレイヤーは TPU をサポートしていないので、文字列演算は TPU では実行されません。(2 つの例外は、tf.keras.layers.Normalization と tf.keras.layers.Rescaling です。これらは TPU で正常に動作し、一般的に画像モデルの最初のレイヤーとして使用されます。)

Model.fit で前処理する

Keras Model.fit を使用する場合、tf.distribute.Strategy.experimental_distribute_dataset や tf.distribute.Strategy.distribute_datasets_from_function でデータを分散する必要はありません。詳細については、前処理レイヤーの使用ガイドと Keras を使用した分散トレーニングガイドを参照してください。次に簡単な例を示します。

strategy = tf.distribute.MirroredStrategy()
with strategy.scope():
  # Create the layer(s) under scope.
  integer_preprocessing_layer = tf.keras.layers.IntegerLookup(vocabulary=FILE_PATH)
  model = ...
  model.compile(...)
dataset = dataset.map(lambda x, y: (integer_preprocessing_layer(x), y))
model.fit(dataset)

Model.fit API を使用する tf.distribute.experimental.ParameterServerStrategy のユーザーは、tf.keras.utils.experimental.DatasetCreator を入力として使用する必要があります。(詳細については、パラメータサーバートレーニングガイドを参照してください)

strategy = tf.distribute.experimental.ParameterServerStrategy(
    cluster_resolver,
    variable_partitioner=variable_partitioner)

with strategy.scope():
  preprocessing_layer = tf.keras.layers.StringLookup(vocabulary=FILE_PATH)
  model = ...
  model.compile(...)

def dataset_fn(input_context):
  ...
  dataset = dataset.map(preprocessing_layer)
  ...
  return dataset

dataset_creator = tf.keras.utils.experimental.DatasetCreator(dataset_fn)
model.fit(dataset_creator, epochs=5, steps_per_epoch=20, callbacks=callbacks)

カスタムトレーニングループを使った前処理

カスタムトレーニングループを作成する場合、tf.distribute.Strategy.experimental_distribute_dataset API または tf.distribute.Strategy.distribute_datasets_from_function API のいずれかを使用してデータを分散します。tf.distribute.Strategy.experimental_distribute_dataset を介してデータセットを分散する場合、これらの前処理 API をデータパイプラインに適用すると、リソースが自動的にデータパイプラインと同じ場所に配置され、リモートリソースアクセスを回避できます。したがって、ここでの例はすべて tf.distribute.Strategy.distribute_datasets_from_function を使用します。この場合、これらの API の初期化を strategy.scope() の下に配置して効率化することが重要です。

strategy = tf.distribute.MirroredStrategy()
vocab = ["a", "b", "c", "d", "f"]

with strategy.scope():
  # Create the layer(s) under scope.
  layer = tf.keras.layers.StringLookup(vocabulary=vocab)

def dataset_fn(input_context):
  # a tf.data.Dataset
  dataset = tf.data.Dataset.from_tensor_slices(["a", "c", "e"]).repeat()

  # Custom your batching, sharding, prefetching, etc.
  global_batch_size = 4
  batch_size = input_context.get_per_replica_batch_size(global_batch_size)
  dataset = dataset.batch(batch_size)
  dataset = dataset.shard(
      input_context.num_input_pipelines,
      input_context.input_pipeline_id)

  # Apply the preprocessing layer(s) to the tf.data.Dataset
  def preprocess_with_kpl(input):
    return layer(input)

  processed_ds = dataset.map(preprocess_with_kpl)
  return processed_ds

distributed_dataset = strategy.distribute_datasets_from_function(dataset_fn)

# Print out a few example batches.
distributed_dataset_iterator = iter(distributed_dataset)
for _ in range(3):
  print(next(distributed_dataset_iterator))

tf.distribute.experimental.ParameterServerStrategy でトレーニングしている場合は、tf.distribute.experimental.coordinator.ClusterCoordinator.create_per_worker_dataset も呼び出すことに注意してください。

@tf.function
def per_worker_dataset_fn():
  return strategy.distribute_datasets_from_function(dataset_fn)

per_worker_dataset = coordinator.create_per_worker_dataset(per_worker_dataset_fn)
per_worker_iterator = iter(per_worker_dataset)

Tensorflow Transform の場合、前述のように、分析段階はトレーニングとは別に行われるため、ここでは触れません。詳細については、チュートリアルを参照してください。通常、この段階では、tf.Transform 前処理関数を作成し、この前処理関数を使用して Apache Beam パイプラインでデータを変換します。分析段階の最後に、トレーニングとサービングの両方に使用できる TensorFlow グラフとして出力をエクスポートできます。この例では、トレーニングパイプラインの部分のみを扱います。

with strategy.scope():
  # working_dir contains the tf.Transform output.
  tf_transform_output = tft.TFTransformOutput(working_dir)
  # Loading from working_dir to create a Keras layer for applying the tf.Transform output to data
  tft_layer = tf_transform_output.transform_features_layer()
  ...

def dataset_fn(input_context):
  ...
  dataset.map(tft_layer, num_parallel_calls=tf.data.AUTOTUNE)
  ...
  return dataset

distributed_dataset = strategy.distribute_datasets_from_function(dataset_fn)

部分バッチ

1）ユーザーが作成する tf.data.Dataset インスタンスに、レプリカの数で均等に除算できないバッチサイズが含まれていない場合、または 2）データセットインスタンスのカーディナリティがバッチサイズで除算できない場合に、部分バッチが発生します。つまり、データセットが複数のレプリカに分散される場合、一部のイテレータでの next 呼び出しが、tf.errors.OutOfRangeError になります。このユースケースに対応するために、tf.distribute は、処理するデータが残っていないレプリカで、バッチサイズ 0 のダミーバッチを返します。

単一のワーカーの場合、データがイテレータの next 呼び出しで返されない場合に、バッチサイズ 0 のダミーバッチが作成され、データセットの実際のデータとともに使用されます。部分バッチの場合、データの最後のグローバルバッチには、データのダミーバッチとともに実際のデータが含まれます。データ処理に使用する抑止条件では、レプリカにデータが存在するかどうかを確認するようになっています。データが存在しないレプリカが検出されると、tf.errors.OutOfRangeError エラーが発生します。

一方、複数のワーカーの場合は、クロスレプリカ通信を使用して各ワーカーのデータの存在を表すブール値が集計されます。これは、すべてのワーカーが分散データセットの処理を終了したことを識別するために使用されます。これにはクロスワーカー通信が伴うため、パフォーマンスに何らかの悪影響が及びます。

注意事項

• マルチワーカーセットアップで tf.distribute.Strategy.experimental_distribute_dataset API を使用する場合、ユーザーはファイルから読み取る tf.data.Dataset を渡します。tf.data.experimental.AutoShardPolicy が AUTO または FILE に設定されている場合、ステップごとの実際のバッチサイズがユーザー定義のグローバルバッチサイズより小さくなる可能性があります。これは、ファイルの残りの要素がグローバルバッチサイズより少なくなる場合に発生することがあります。ユーザーは、実行するステップ数に依存せずにデータセットを使い果たすか、tf.data.experimental.AutoShardPolicy を DATA に設定してこれを回避することができます。

• ステートフルデータセット変換は、現在 tf.distribute でサポートされていません。データセットにあるステートフル演算は、現在のところ無視されます。たとえば、データセットに tf.random.uniform を使って画像を回転させる map_fn がある場合、Python プロセスが実行されているローカルマシン上の状態 (ランダムシード) に依存するデータセットグラフがあります。

• デフォルトで無効化されている実験的な tf.data.experimental.OptimizationOptions は、tf.distribute と併せて使用されている場合などでは、パフォーマンスの低下を生じる可能性があります。分散環境においてワークロードのパフォーマンスに有益であることが確認されてから、有効化するようにしてください。

• 一般的に tf.data を使用して入力パイプラインを最適化する方法については、このガイドを参照してください。また、以下のヒントをご覧ください。

  • 複数のワーカーがあり、tf.data.Dataset.list_files を使用して、1 つ以上の glob パターンに一致するすべてのファイルからデータセットを作成している場合は、必ず seed 引数を設定するか、shuffle=False を設定して、各ワーカーが一貫してファイルを分割するようにします。

• 入力パイプラインにレコードレベルでのデータのシャッフルとデータの解析の両方が含まれている場合、解析されていないデータが解析されたデータよりも大幅に大きくない限り (通常はそうではありません)、次の例に示すように、最初にシャッフルしてから解析します。これにより、メモリ使用量とパフォーマンスが向上することがあります。

d = tf.data.Dataset.list_files(pattern, shuffle=False)
d = d.shard(num_workers, worker_index)
d = d.repeat(num_epochs)
d = d.shuffle(shuffle_buffer_size)
d = d.interleave(tf.data.TFRecordDataset,
                 cycle_length=num_readers, block_length=1)
d = d.map(parser_fn, num_parallel_calls=num_map_threads)

• tf.data.Dataset.shuffle(buffer_size, seed=None, reshuffle_each_iteration=None) は、buffer_size 要素の内部バッファを維持し、buffer_size を削減します。これにより、OOM の問題が軽減される可能性があります。

• tf.distribute.experimental_distribute_dataset またはtf.distribute.distribute_datasets_from_function を使用している場合、ワーカーがデータを処理する順番は保証されていません。これは通常、tf.distribute を使用して予測をスケーリングする場合に必要です。ただし、バッチの各要素に対するインデックスを挿入し、それに従って出力を順序付けることができます。次のスニペットは、出力を順序付ける方法を示します。

注意: ここでは便宜上、tf.distribute.MirroredStrategy が使用されていますが、複数のワーカーを使用しており、単一ワーカーへの分散に tf.distribute.MirroredStrategy が使用されている場合には、入力の順番のみを変更する必要があります。

mirrored_strategy = tf.distribute.MirroredStrategy()
dataset_size = 24
batch_size = 6
dataset = tf.data.Dataset.range(dataset_size).enumerate().batch(batch_size)
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)

def predict(index, inputs):
  outputs = 2 * inputs
  return index, outputs

result = {}
for index, inputs in dist_dataset:
  output_index, outputs = mirrored_strategy.run(predict, args=(index, inputs))
  indices = list(mirrored_strategy.experimental_local_results(output_index))
  rindices = []
  for a in indices:
    rindices.extend(a.numpy())
  outputs = list(mirrored_strategy.experimental_local_results(outputs))
  routputs = []
  for a in outputs:
    routputs.extend(a.numpy())
  for i, value in zip(rindices, routputs):
    result[i] = value

print(result)

## tf.data の代わりにテンソル入力を使用する

入力を表現するtf.data.Datasetと、上記に示した、複数のデバイスにデータセットを分散する後続の API を使用できないことがあります。このような場合は、生のテンソルを使用するか、ジェネレータの入力を使用することができます。

任意のテンソル入力に experimental_distribute_values_from_function を使用する

strategy.run は、next(iterator) の出力である tf.distribute.DistributedValues を受け入れます。テンソル値を渡すには、tf.distribute.Strategy.experimental_distribute_values_from_function を使用して生のテンソルから tf.distribute.DistributedValues を構築します。ユーザーは、tf.distribute.experimental.ValueContext 入力オブジェクトを使用して、このオプションを使用して入力関数で独自のバッチ処理およびシャーディングロジックを指定する必要があります。

mirrored_strategy = tf.distribute.MirroredStrategy()

def value_fn(ctx):
  return tf.constant(ctx.replica_id_in_sync_group)

distributed_values = mirrored_strategy.experimental_distribute_values_from_function(value_fn)
for _ in range(4):
  result = mirrored_strategy.run(lambda x: x, args=(distributed_values,))
  print(result)

ジェネレータからの入力である場合に tf.data.Dataset.from_generator を使用する

使用を検討しているジェネレータ関数がある場合は、from_generator API を使用してtf.data.Dataset インスタンスを作成できます。

注意: 現在のところ、tf.distribute.TPUStrategy ではサポートされていません。

mirrored_strategy = tf.distribute.MirroredStrategy()
def input_gen():
  while True:
    yield np.random.rand(4)

# use Dataset.from_generator
dataset = tf.data.Dataset.from_generator(
    input_gen, output_types=(tf.float32), output_shapes=tf.TensorShape([4]))
dist_dataset = mirrored_strategy.experimental_distribute_dataset(dataset)
iterator = iter(dist_dataset)
for _ in range(4):
  result = mirrored_strategy.run(lambda x: x, args=(next(iterator),))
  print(result)