Python tf.data.experimental.service.distribute用法及代码示例

将数据集处理移动到 tf.data 服务的转换。

用法

tf.data.experimental.service.distribute(
    processing_mode, service, job_name=None, consumer_index=None,
    num_consumers=None, max_outstanding_requests=None, data_transfer_protocol=None,
    compression='AUTO', target_workers='AUTO'
)

当您迭代包含 distribute 转换的数据集时，tf.data 服务会创建一个 "job"，它会为数据集迭代生成数据。

tf.data 服务使用一组工作人员来准备数据以训练您的模型。 tf.data.experimental.service.distribute 的 processing_mode 参数说明了如何利用多个工作人员来处理输入数据集。目前有两种处理模式可供选择："distributed_epoch"和"parallel_epochs"。

"distributed_epoch" 表示数据集将被拆分到所有 tf​​.data 服务工作者中。调度程序为数据集生成"splits"，并将它们发送给工作人员进行进一步处理。例如，如果数据集以文件名列表开头，则调度程序将遍历文件名并将文件名发送给 tf.data 工作人员，后者将对这些文件执行其余的数据集转换。 "distributed_epoch" 当您的模型需要仅查看数据集的每个元素一次，或者需要按 generally-sequential 顺序查看数据时很有用。 "distributed_epoch" 仅适用于具有可拆分源的数据集，例如 Dataset.from_tensor_slices 、 Dataset.list_files 或 Dataset.range 。

"parallel_epochs" 表示整个输入数据集将由每个 tf.data 服务工作者独立处理。出于这个原因，重要的是非确定性地打乱数据(例如文件名)，以便每个工作人员以不同的顺序处理数据集的元素。 "parallel_epochs" 可用于分发不可拆分的数据集。

如果有两个工作人员，"parallel_epochs" 将生成数据集的每个元素两次：

dispatcher = tf.data.experimental.service.DispatchServer()
dispatcher_address = dispatcher.target.split("://")[1]
# Start two workers
workers = [
    tf.data.experimental.service.WorkerServer(
        tf.data.experimental.service.WorkerConfig(
            dispatcher_address=dispatcher_address)) for _ in range(2)
]
dataset = tf.data.Dataset.range(10)
dataset = dataset.apply(tf.data.experimental.service.distribute(
    processing_mode="parallel_epochs", service=dispatcher.target))
print(sorted(list(dataset.as_numpy_iterator())))
[0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9]

另一方面，"distributed_epoch" 仍然会生成每个元素一次：

dispatcher = tf.data.experimental.service.DispatchServer()
dispatcher_address = dispatcher.target.split("://")[1]
workers = [
    tf.data.experimental.service.WorkerServer(
        tf.data.experimental.service.WorkerConfig(
            dispatcher_address=dispatcher_address)) for _ in range(2)
]
dataset = tf.data.Dataset.range(10)
dataset = dataset.apply(tf.data.experimental.service.distribute(
    processing_mode="distributed_epoch", service=dispatcher.target))
print(sorted(list(dataset.as_numpy_iterator())))
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

使用 apply(tf.data.experimental.service.distribute(...)) 时，apply 转换之前的数据集在 tf.data 服务中执行，而 apply 之后的操作发生在本地进程中。

dispatcher = tf.data.experimental.service.DispatchServer()
dispatcher_address = dispatcher.target.split("://")[1]
workers = [
    tf.data.experimental.service.WorkerServer(
        tf.data.experimental.service.WorkerConfig(
            dispatcher_address=dispatcher_address)) for _ in range(2)
]
dataset = tf.data.Dataset.range(5)
dataset = dataset.map(lambda x:x*x)
dataset = dataset.apply(
   tf.data.experimental.service.distribute("parallel_epochs",
                                           dispatcher.target))
dataset = dataset.map(lambda x:x+1)
print(sorted(list(dataset.as_numpy_iterator())))
[1, 1, 2, 2, 5, 5, 10, 10, 17, 17]

在上面的示例中，数据集操作(在对元素应用distribute 函数之前)将在 tf.data 工作人员上执行，并且通过 RPC 提供元素。剩余的转换(在调用 distribute 之后)将在本地执行。调度程序和工作人员将绑定到使用的空闲端口(随机选择)，以便相互通信。但是，要将它们绑定到特定端口，可以传递 port 参数。

job_name 参数允许在多个数据集之间共享作业。不是每个数据集都创建自己的作业，而是所有具有相同 job_name 的数据集将从同一个作业中使用。将为数据集的每次迭代创建一个新作业(每次重复 Dataset.repeat 计为一次新迭代)。假设 DispatchServer 在 localhost:5000 上提供服务，并且两个训练工作者(在单个客户端或 multi-client 设置中)迭代以下数据集，并且有一个 tf.data 工作者：

range5_dataset = tf.data.Dataset.range(5)
dataset = range5_dataset.apply(tf.data.experimental.service.distribute(
    "parallel_epochs", "localhost:5000", job_name="my_job_name"))
for iteration in range(3):
  print(list(dataset))

每个作业的元素将在两个进程之间拆分，进程在 first-come first-served 的基础上使用元素。一种可能的结果是进程 1 打印

[0, 2, 4]
[0, 1, 3]
[1]

并处理 2 个打印件

[1, 3]
[2, 4]
[0, 2, 3, 4]

在 tf.data 服务的生命周期内，不得在不同的训练作业中重复使用作业名称。一般来说，tf.data 服务预计会在单个训练作业的持续时间内存在。要将 tf.data 服务用于多个训练作业，请确保使用不同的作业名称以避免冲突。例如，假设训练作业使用 job_name="job" 调用 distribute 并读取直到输入结束。如果另一个独立作业连接到同一个 tf.data 服务并尝试从 job_name="job" 读取，它将立即接收输入结束，而不会获取任何数据。

协调数据读取

默认情况下，当多个消费者从同一个作业中读取数据时，他们会在 first-come first-served 的基础上接收数据。在某些用例中，协调消费者是有利的。在每一步，消费者都从同一个工作人员那里读取数据。

例如，tf.data 服务可用于在同步训练期间跨集群协调示例大小，以便在每个步骤中，所有副本都在 similar-sized 元素上进行训练。为此，请定义一个数据集，该数据集生成几轮 num_consumers 连续 similar-sized 批次，然后通过设置 consumer_index 和 num_consumers 启用协调读取。

注意：为了让消费者保持同步，循环数据消费要求数据集具有无限的基数。您可以通过在数据集定义的末尾添加 .repeat() 来获得此信息。

Keras 和分发策略

distribute 转换生成的数据集可以像任何其他 tf.data.Dataset 一样传递到 Keras 的 Model.fit 或 Distribution Strategy 的 tf.distribute.Strategy.experimental_distribute_dataset 。我们建议在调用 distribute 时设置 job_name，这样如果有多个工作人员，他们就可以从同一个作业中读取数据。请注意，在设置 job_name 时，通常由 experimental_distribute_dataset 执行的自动分片将被禁用，因为共享作业已经导致在工作人员之间拆分数据。当使用共享作业时，数据将在工作人员之间动态平衡，以便他们大约在同一时间到达输入结束。与自动分片相比，这会导致更好的工作人员利用率，其中每个工作人员处理一组独立的文件，并且一些工作人员可能比其他工作人员更早地用完数据。