Python tf.distribute.MirroredStrategy用法及代码示例

在一台机器上跨多个副本进行同步训练。

继承自：Strategy

用法

tf.distribute.MirroredStrategy(
    devices=None, cross_device_ops=None
)

此策略通常用于在具有多个 GPU 的一台机器上进行训练。对于 TPU，请使用 tf.distribute.TPUStrategy 。要将MirroredStrategy 与多个工人一起使用，请参阅tf.distribute.experimental.MultiWorkerMirroredStrategy。

例如，在 MirroredStrategy 下创建的变量是 MirroredVariable 。如果在策略的构造函数参数中没有指定设备，那么它将使用所有可用的 GPU。如果没有找到 GPU，它将使用可用的 CPU。请注意，TensorFlow 将机器上的所有 CPU 视为单个设备，并在内部使用线程进行并行处理。

strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
with strategy.scope():
  x = tf.Variable(1.)
x
MirroredVariable:{
  0:<tf.Variable ... shape=() dtype=float32, numpy=1.0>,
  1:<tf.Variable ... shape=() dtype=float32, numpy=1.0>
}

在使用分配策略时，所有变量的创建都应该在策略的范围内完成。这将在所有副本之间复制变量，并使用all-reduce 算法使它们保持同步。

在用 tf.function 包装的 MirroredStrategy 中创建的变量仍然是 MirroredVariables 。

x = []
@tf.function  # Wrap the function with tf.function.
def create_variable():
  if not x:
    x.append(tf.Variable(1.))
  return x[0]
strategy = tf.distribute.MirroredStrategy(["GPU:0", "GPU:1"])
with strategy.scope():
  _ = create_variable()
  print(x[0])
MirroredVariable:{
  0:<tf.Variable ... shape=() dtype=float32, numpy=1.0>,
  1:<tf.Variable ... shape=() dtype=float32, numpy=1.0>
}

experimental_distribute_dataset 可用于在编写自己的训练循环时将数据集分布在副本之间。如果您使用 tf.keras 中可用的 .fit 和 .compile 方法，那么 tf.keras 将为您处理分发。

例如：

my_strategy = tf.distribute.MirroredStrategy()
with my_strategy.scope():
  @tf.function
  def distribute_train_epoch(dataset):
    def replica_fn(input):
      # process input and return result
      return result

    total_result = 0
    for x in dataset:
      per_replica_result = my_strategy.run(replica_fn, args=(x,))
      total_result += my_strategy.reduce(tf.distribute.ReduceOp.SUM,
                                         per_replica_result, axis=None)
    return total_result

  dist_dataset = my_strategy.experimental_distribute_dataset(dataset)
  for _ in range(EPOCHS):
    train_result = distribute_train_epoch(dist_dataset)