在 TensorFlow.org 上查看 | 在 Google Colab 中运行 | 在 GitHub 上查看源代码 | 下载笔记本 |
该 tf.distribute API 为用户提供了一种简单的方法,可以将他们的训练从单台机器扩展到多台机器。在扩展模型时,用户还必须将输入分布到多个设备上。 tf.distribute
提供 API,您可以使用这些 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
之外的其他方法,请参考本指南中的 张量输入部分。在非分布式训练循环中,首先创建一个 tf.data.Dataset
实例,然后迭代其元素。例如
import tensorflow as tf
# Helper libraries
import numpy as np
import os
print(tf.__version__)
2023-12-07 02:57:52.101761: 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 2023-12-07 02:57:52.101810: 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 2023-12-07 02:57:52.103319: 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
策略,引入了两个 API,它们将分布一个 tf.data.Dataset
实例并返回一个分布式数据集对象。然后,用户可以迭代此分布式数据集实例并像以前一样训练他们的模型。现在让我们更详细地了解这两个 API - tf.distribute.Strategy.experimental_distribute_dataset
和 tf.distribute.Strategy.distribute_datasets_from_function
tf.distribute.Strategy.experimental_distribute_dataset
用法
此 API 以 tf.data.Dataset
实例作为输入,并返回一个 tf.distribute.DistributedDataset
实例。您应该使用等于全局批处理大小的值对输入数据集进行批处理。此全局批处理大小是您希望在 1 步中跨所有设备处理的样本数量。您可以以 Pythonic 方式迭代此分布式数据集,或者使用 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
实例,该批处理大小等于全局批处理大小除以同步副本的数量。同步副本的数量等于参与训练期间梯度全减少的设备数量。当用户在分布式迭代器上调用 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
指定。请注意,在使用 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
设置三个不同的选项
- AUTO: 这是默认选项,这意味着将尝试按 FILE 进行分片。如果未检测到基于文件的数据集,则按 FILE 进行分片的尝试将失败。
tf.distribute
然后将回退到按 DATA 进行分片。请注意,如果输入数据集是基于文件的,但文件数量少于工作器数量,则会引发InvalidArgumentError
。如果发生这种情况,请显式将策略设置为AutoShardPolicy.DATA
,或将您的输入源拆分为更小的文件,使文件数量大于工作器数量。 FILE: 这是您希望在所有工作器上分片输入文件时的选项。如果您输入文件数量远大于工作器数量,并且文件中的数据分布均匀,则应使用此选项。此选项的缺点是,如果文件中的数据分布不均匀,则工作器会处于空闲状态。如果文件数量少于工作器数量,则会引发
InvalidArgumentError
。如果发生这种情况,请显式将策略设置为AutoShardPolicy.DATA
。例如,让我们在 2 个工作器上分布 2 个文件,每个工作器 1 个副本。文件 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: 这将在所有工作器上自动分片元素。每个工作器将读取整个数据集,并且只处理分配给它的分片。所有其他分片将被丢弃。这通常用于输入文件数量少于工作器数量,并且您希望在所有工作器上更好地分片数据。缺点是每个工作器都会读取整个数据集。例如,让我们在 2 个工作器上分布 1 个文件。文件 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: 如果您关闭自动分片,每个工作器将处理所有数据。例如,让我们在 2 个工作器上分布 1 个文件。文件 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 设备上调用输入函数。除了允许用户指定自己的批处理和分片逻辑外,与用于多工作节点训练的 tf.distribute.Strategy.experimental_distribute_dataset
相比,此 API 还展示了更好的可扩展性和性能。
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.data.Dataset
实例类似,您需要在 tf.distribute.DistributedDataset
实例上创建一个迭代器,以迭代它并访问 tf.distribute.DistributedDataset
中的元素。以下是可以创建 tf.distribute.DistributedIterator
并使用它来训练模型的方法
用法
使用 Pythonic for 循环结构
您可以使用用户友好的 Pythonic 循环来迭代 tf.distribute.DistributedDataset
。从 tf.distribute.DistributedIterator
返回的元素可以是单个 tf.Tensor
或 tf.distribute.DistributedValues
,其中包含每个副本的值。将循环放在 tf.function
中将提高性能。但是,目前不支持对 tf.distribute.DistributedDataset
的循环使用 break
和 return
,该循环位于 tf.function
中。
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()
。前两者本质上是相同的
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
。但是,在数据准备好用于模型之前,需要对其进行预处理,例如通过清理、转换和增强。两组这些方便的工具是
Keras 预处理层:一组 Keras 层,允许开发人员构建 Keras 原生的输入处理管道。一些 Keras 预处理层包含不可训练的状态,这些状态可以在初始化时设置或
adapt
(请参阅 Keras 预处理层指南 的adapt
部分)。在分布有状态的预处理层时,应将状态复制到所有工作节点。要使用这些层,您可以将它们作为模型的一部分,也可以将它们应用于数据集。TensorFlow Transform (tf.Transform):一个用于 TensorFlow 的库,允许您通过数据预处理管道定义实例级和全通数据转换。Tensorflow Transform 包含两个阶段。第一个是分析阶段,在该阶段,原始训练数据在全通过程中进行分析,以计算转换所需的统计信息,并将转换逻辑生成为实例级操作。第二个是转换阶段,在该阶段,原始训练数据在实例级过程中进行转换。
Keras 预处理层与 Tensorflow Transform
Tensorflow Transform 和 Keras 预处理层都提供了一种方法,可以在训练期间将预处理分离出来,并在推理期间将预处理与模型捆绑在一起,从而减少训练/服务偏差。
Tensorflow Transform 与 TFX 深度集成,提供了一个可扩展的 map-reduce 解决方案,可以在与训练管道不同的作业中分析和转换任何大小的数据集。如果您需要对无法容纳在单个机器上的数据集运行分析,则 Tensorflow Transform 应该是您的首选。
Keras 预处理层更侧重于在训练期间,从磁盘读取数据后应用的预处理。它们与 Keras 库中的模型开发无缝衔接。它们支持通过 adapt
分析较小的数据集,并支持图像数据增强等用例,其中对输入数据集的每次遍历都会产生不同的训练示例。
这两个库也可以混合使用,其中 Tensorflow Transform 用于分析和对输入数据进行静态转换,而 Keras 预处理层用于训练时转换(例如,独热编码或数据增强)。
使用 tf.distribute 的最佳实践
使用这两个工具都涉及初始化要应用于数据的转换逻辑,这可能会创建 Tensorflow 资源。这些资源或状态应复制到所有工作器,以节省工作器间或工作器-协调器之间的通信。为此,建议您在 tf.distribute.Strategy.scope
下创建 Keras 预处理层、tft.TFTransformOutput.transform_features_layer
或 tft.TransformFeaturesLayer
,就像您对任何其他 Keras 层所做的那样。
以下示例演示了使用 tf.distribute.Strategy
API 与高级 Keras Model.fit
API 以及自定义训练循环分别使用。
针对 Keras 预处理层用户的额外说明
预处理层和大型词汇表
在多工作器设置(例如,tf.distribute.MultiWorkerMirroredStrategy
、tf.distribute.experimental.ParameterServerStrategy
、tf.distribute.TPUStrategy
)中处理大型词汇表(超过 1 GB)时,建议将词汇表保存到所有工作器都可以访问的静态文件(例如,使用 Cloud Storage)。这将减少在训练期间将词汇表复制到所有工作器所花费的时间。
在 tf.data
管道中进行预处理与在模型中进行预处理
虽然 Keras 预处理层可以作为模型的一部分或直接应用于 tf.data.Dataset
,但每个选项都有其优势。
- 在模型中应用预处理层使您的模型可移植,并有助于减少训练/服务偏差。(有关更多详细信息,请参阅 使用预处理层指南 中的“在推理时在模型内部进行预处理的好处”部分。)
- 在
tf.data
管道中应用预处理层允许预取或卸载到 CPU,这在使用加速器时通常会提供更好的性能。
在运行一个或多个 TPU 时,用户几乎总是应该将 Keras 预处理层放在 tf.data
管道中,因为并非所有层都支持 TPU,并且字符串操作不会在 TPU 上执行。(两个例外是 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)
使用 tf.distribute.experimental.ParameterServerStrategy
和 Model.fit
API 的用户需要使用 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
从匹配一个或多个 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
来扩展预测,通常需要这样做。但是,您可以为批次中的每个元素插入索引,并相应地对输出进行排序。以下代码片段演示了如何对输出进行排序。
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
有时用户无法使用 tf.data.Dataset
来表示他们的输入,因此无法使用上述 API 将数据集分配到多个设备。在这种情况下,您可以使用原始张量或来自生成器的输入。
对任意张量输入使用 experimental_distribute_values_from_function
strategy.run
接受 tf.distribute.DistributedValues
,它是 next(iterator)
的输出。要传递张量值,请使用 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
实例。
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.795073], shape=(1,), dtype=float32), 1: tf.Tensor([0.4941732], shape=(1,), dtype=float32), 2: tf.Tensor([0.51117146], shape=(1,), dtype=float32), 3: tf.Tensor([0.791901], shape=(1,), dtype=float32) } PerReplica:{ 0: tf.Tensor([0.10990978], shape=(1,), dtype=float32), 1: tf.Tensor([0.61591166], shape=(1,), dtype=float32), 2: tf.Tensor([0.17349982], shape=(1,), dtype=float32), 3: tf.Tensor([0.8937937], shape=(1,), dtype=float32) } PerReplica:{ 0: tf.Tensor([0.97211426], shape=(1,), dtype=float32), 1: tf.Tensor([0.30425492], shape=(1,), dtype=float32), 2: tf.Tensor([0.80144566], shape=(1,), dtype=float32), 3: tf.Tensor([0.25493157], shape=(1,), dtype=float32) } PerReplica:{ 0: tf.Tensor([0.07450782], shape=(1,), dtype=float32), 1: tf.Tensor([0.23319475], shape=(1,), dtype=float32), 2: tf.Tensor([0.22552523], shape=(1,), dtype=float32), 3: tf.Tensor([0.7449827], shape=(1,), dtype=float32) }