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 執行的自動分片將被禁用，因為共享作業已經導致在工作人員之間拆分數據。當使用共享作業時，數據將在工作人員之間動態平衡，以便他們大約在同一時間到達輸入結束。與自動分片相比，這會導致更好的工作人員利用率，其中每個工作人員處理一組獨立的文件，並且一些工作人員可能比其他工作人員更早地用完數據。