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

TPU 和 TPU Pod 的同步训练。

继承自：Strategy

用法

tf.distribute.TPUStrategy(
    tpu_cluster_resolver=None, experimental_device_assignment=None,
    experimental_spmd_xla_partitioning=False
)

要构造 TPUStrategy 对象，您需要运行如下初始化代码：

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
strategy = tf.distribute.TPUStrategy(resolver)

在使用分发策略时，在策略范围内创建的变量将在所有副本中复制，并且可以使用 all-reduce 算法保持同步。

要在 TPU 上运行 TF2 程序，您可以将 tf.keras 中的 .compile 和 .fit API 与 TPUStrategy 一起使用，或者通过直接调用 strategy.run 来编写自己的自定义训练循环。请注意，TPUStrategy 不支持纯粹的 Eager 执行，因此请确保传递给 strategy.run 的函数是 tf.function 或 strategy.run 如果启用了 Eager 行为，则在 tf.function 内调用。在 https://www.tensorflow.org/guide/tpu 中查看更多详细信息。

在编写自己的训练循环时，distribute_datasets_from_function 和 experimental_distribute_dataset API 可用于在 TPU 工作人员之间分配数据集。如果您使用 tf.keras.Model 中可用的 fit 和 compile 方法，那么 Keras 将为您处理分发。

在 TPU 上编写自定义训练循环的示例：

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)
iterator = iter(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),))

train_step(iterator)

对于模型并行等高级用例，您可以在创建 TPUStrategy 时设置 experimental_device_assignment 参数以指定副本数和逻辑设备数。下面是一个使用 2 个逻辑设备和 1 个副本初始化 TPU 系统的示例。

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
topology = tf.tpu.experimental.initialize_tpu_system(resolver)
device_assignment = tf.tpu.experimental.DeviceAssignment.build(
    topology,
    computation_shape=[1, 1, 1, 2],
    num_replicas=1)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)

然后您可以仅在逻辑设备 0 上运行 tf.add 操作。

@tf.function()
def step_fn(inputs):
  features, _ = inputs
  output = tf.add(features, features)

  # Add operation will be executed on logical device 0.
  output = strategy.experimental_assign_to_logical_device(output, 0)
  return output
dist_dataset = strategy.distribute_datasets_from_function(
    dataset_fn)
iterator = iter(dist_dataset)
strategy.run(step_fn, args=(next(iterator),))

experimental_spmd_xla_partitioning 启用实验性 XLA SPMD 函数以实现模型并行性。该标志可以减少编译时间和 HBM 要求。在此模式下运行时，每个输入张量必须被分区(通过 strategy.experimental_split_to_logical_devices )或完全复制(通过 strategy.experimental_replicate_to_logical_devices )到所有逻辑设备。在此模式下调用 strategy.experimental_assign_to_logical_device 将导致 ValueError。