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

一种对多个工人进行同步训练的分配策略。

继承自：MultiWorkerMirroredStrategy，Strategy

用法

tf.distribute.experimental.MultiWorkerMirroredStrategy(
    communication=tf.distribute.experimental.CollectiveCommunication.AUTO,
    cluster_resolver=None
)

该策略实现了跨多个工作人员的同步分布式训练，每个工作人员都可能具有多个 GPU。与 tf.distribute.MirroredStrategy 类似，它将所有变量和计算复制到每个本地设备。不同之处在于它使用分布式集体实现(例如all-reduce)，以便多个工作人员可以一起工作。

您需要在每个工作人员上启动程序并正确配置cluster_resolver。例如，如果您使用 tf.distribute.cluster_resolver.TFConfigClusterResolver ，则每个工作人员都需要在 TF_CONFIG 环境变量中设置其对应的 task_type 和 task_id。两个工作集群的 worker-0 上的示例 TF_CONFIG 是：

TF_CONFIG = '{"cluster":{"worker":["localhost:12345", "localhost:23456"]}, "task":{"type":"worker", "index":0} }'

您的程序在每个工人as-is 上运行。请注意，集体要求每个工人都参与。所有tf.distribute 和非tf.distribute API 都可以在内部使用集合，例如检查点和保存，因为读取带有 tf.VariableSynchronization.ON_READ all-reduces 值的 tf.Variable。因此，建议在每个工人上运行完全相同的程序。根据worker的task_type或task_id调度为error-prone。

cluster_resolver.num_accelerators() 确定策略使用的 GPU 数量。如果为零，则该策略使用 CPU。所有工作人员都需要使用相同数量的设备，否则行为未定义。

此策略不适用于 TPU。请改用tf.distribute.TPUStrategy。

设置 TF_CONFIG 后，使用此策略类似于使用 tf.distribute.MirroredStrategy 和 tf.distribute.TPUStrategy 。

strategy = tf.distribute.MultiWorkerMirroredStrategy()

with strategy.scope():
  model = tf.keras.Sequential([
    tf.keras.layers.Dense(2, input_shape=(5,)),
  ])
  optimizer = tf.keras.optimizers.SGD(learning_rate=0.1)

def dataset_fn(ctx):
  x = np.random.random((2, 5)).astype(np.float32)
  y = np.random.randint(2, size=(2, 1))
  dataset = tf.data.Dataset.from_tensor_slices((x, y))
  return dataset.repeat().batch(1, drop_remainder=True)
dist_dataset = strategy.distribute_datasets_from_function(dataset_fn)

model.compile()
model.fit(dist_dataset)

您还可以编写自己的训练循环：

@tf.function
def train_step(iterator):

  def step_fn(inputs):
    features, labels = inputs
    with tf.GradientTape() as tape:
      logits = model(features, training=True)
      loss = tf.keras.losses.sparse_categorical_crossentropy(
          labels, logits)

    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

  strategy.run(step_fn, args=(next(iterator),))

for _ in range(NUM_STEP):
  train_step(iterator)

有关详细教程，请参阅使用 Keras 进行 Multi-worker 训练。

保存

您需要在所有工作人员上保存和检查点，而不仅仅是一个。这是因为同步=ON_READ的变量在保存期间会触发聚合。建议在每个工作人员上保存到不同的路径以避免竞争条件。每个工人保存相同的东西。有关示例，请参阅 Multi-worker 使用 Keras 教程进行训练。

已知的问题

• tf.distribute.cluster_resolver.TFConfigClusterResolver 未返回正确数量的加速器。如果 cluster_resolver 是 tf.distribute.cluster_resolver.TFConfigClusterResolver 或 None ，则该策略使用所有可用的 GPU。
• 在 Eager 模式下，需要在调用任何其他 Tensorflow API 之前创建策略。