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__)
2024-01-11 18:25:57.072759: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered 2024-01-11 18:25:57.072803: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered 2024-01-11 18:25:57.074336: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered 2.15.0
# 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))
Available devices: 0) LogicalDevice(name='/device:CPU:0', device_type='CPU') 1) LogicalDevice(name='/device:CPU:1', device_type='CPU') 2) LogicalDevice(name='/device:GPU:0', device_type='GPU') 3) LogicalDevice(name='/device:GPU:1', device_type='GPU') 4) LogicalDevice(name='/device:GPU:2', device_type='GPU') 5) LogicalDevice(name='/device:GPU:3', device_type='GPU')
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.Tensor( [[0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7]], shape=(16, 1), dtype=float32) tf.Tensor( [[0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7]], shape=(16, 1), dtype=float32) tf.Tensor( [[0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7]], shape=(16, 1), dtype=float32) tf.Tensor( [[0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7]], shape=(16, 1), dtype=float32) tf.Tensor( [[0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7]], shape=(16, 1), dtype=float32) tf.Tensor( [[0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7] [0.7]], shape=(16, 1), dtype=float32) tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32)
ユーザーの既存のコードへの変更を最小限に抑えて 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)))
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3') (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> })
プロパティ
バッチ処理
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: []
- バッチ 1:
- 分散無し:
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: []
- バッチ 1:
- 分散無し:
注意: 上記の例は、異なるレプリカでグローバルバッチがどのように分割されるかのみを説明しています。実装によって実際の値が異なる可能性があるため、各レプリカで最終的に得られる可能性のある実際の値に依存することはお勧めできません。
データセットのバッチの再作成には、レプリカの数とともに直線的に増加する空間的コストがあります。つまり、マルチワーカートレーニングのユースケースで言えば、入力パイプラインで 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]
- ワーカー 0:
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]
- ワーカー 0:
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)
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3')
プロパティ
バッチ処理
入力関数の戻り値である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)
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3') Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor([[0.7]], shape=(1, 1), dtype=float32), 1: tf.Tensor([[0.7]], shape=(1, 1), dtype=float32), 2: tf.Tensor([[0.7]], shape=(1, 1), dtype=float32), 3: tf.Tensor([[0.7]], shape=(1, 1), dtype=float32) }
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)
Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) } Loss is PerReplica:{ 0: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 1: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 2: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32), 3: tf.Tensor( [[0.7] [0.7] [0.7] [0.7]], shape=(4, 1), dtype=float32) }
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)
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3') ([0], [1], [2], [3]) ([4], [5], [6], [7]) ([8], [], [], [])
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)
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3') (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }) (PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> }, PerReplica:{ 0: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 1: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 2: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)>, 3: <tf.Tensor: shape=(4, 1), dtype=float32, numpy= array([[1.], [1.], [1.], [1.]], dtype=float32)> })
データの前処理をする
以上、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))
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3') PerReplica:{ 0: tf.Tensor([1], shape=(1,), dtype=int64), 1: tf.Tensor([3], shape=(1,), dtype=int64), 2: tf.Tensor([0], shape=(1,), dtype=int64), 3: tf.Tensor([1], shape=(1,), dtype=int64) } PerReplica:{ 0: tf.Tensor([3], shape=(1,), dtype=int64), 1: tf.Tensor([0], shape=(1,), dtype=int64), 2: tf.Tensor([1], shape=(1,), dtype=int64), 3: tf.Tensor([3], shape=(1,), dtype=int64) } PerReplica:{ 0: tf.Tensor([0], shape=(1,), dtype=int64), 1: tf.Tensor([1], shape=(1,), dtype=int64), 2: tf.Tensor([3], shape=(1,), dtype=int64), 3: tf.Tensor([0], shape=(1,), dtype=int64) }
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)
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3') WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance. WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance. WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance. WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance. {0: 0, 1: 2, 2: 4, 3: 6, 4: 8, 5: 10, 6: 12, 7: 14, 8: 16, 9: 18, 10: 20, 11: 22, 12: 24, 13: 26, 14: 28, 15: 30, 16: 32, 17: 34, 18: 36, 19: 38, 20: 40, 21: 42, 22: 44, 23: 46}
入力を表現する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)
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3') WARNING:tensorflow:Using MirroredStrategy eagerly has significant overhead currently. We will be working on improving this in the future, but for now please wrap `call_for_each_replica` or `experimental_run` or `run` inside a tf.function to get the best performance. PerReplica:{ 0: tf.Tensor(0, shape=(), dtype=int32), 1: tf.Tensor(1, shape=(), dtype=int32), 2: tf.Tensor(2, shape=(), dtype=int32), 3: tf.Tensor(3, shape=(), dtype=int32) } PerReplica:{ 0: tf.Tensor(0, shape=(), dtype=int32), 1: tf.Tensor(1, shape=(), dtype=int32), 2: tf.Tensor(2, shape=(), dtype=int32), 3: tf.Tensor(3, shape=(), dtype=int32) } PerReplica:{ 0: tf.Tensor(0, shape=(), dtype=int32), 1: tf.Tensor(1, shape=(), dtype=int32), 2: tf.Tensor(2, shape=(), dtype=int32), 3: tf.Tensor(3, shape=(), dtype=int32) } PerReplica:{ 0: tf.Tensor(0, shape=(), dtype=int32), 1: tf.Tensor(1, shape=(), dtype=int32), 2: tf.Tensor(2, shape=(), dtype=int32), 3: tf.Tensor(3, shape=(), dtype=int32) }
ジェネレータからの入力である場合に 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)
INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3') PerReplica:{ 0: tf.Tensor([0.40325633], shape=(1,), dtype=float32), 1: tf.Tensor([0.28778756], shape=(1,), dtype=float32), 2: tf.Tensor([0.5146715], shape=(1,), dtype=float32), 3: tf.Tensor([0.3218396], shape=(1,), dtype=float32) } PerReplica:{ 0: tf.Tensor([0.79922175], shape=(1,), dtype=float32), 1: tf.Tensor([0.02518538], shape=(1,), dtype=float32), 2: tf.Tensor([0.27494904], shape=(1,), dtype=float32), 3: tf.Tensor([0.54404545], shape=(1,), dtype=float32) } PerReplica:{ 0: tf.Tensor([0.7162087], shape=(1,), dtype=float32), 1: tf.Tensor([0.76167136], shape=(1,), dtype=float32), 2: tf.Tensor([0.8244246], shape=(1,), dtype=float32), 3: tf.Tensor([0.37525535], shape=(1,), dtype=float32) } PerReplica:{ 0: tf.Tensor([0.8093572], shape=(1,), dtype=float32), 1: tf.Tensor([0.0389544], shape=(1,), dtype=float32), 2: tf.Tensor([0.5250396], shape=(1,), dtype=float32), 3: tf.Tensor([0.04613635], shape=(1,), dtype=float32) }