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)