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

用法

experimental_split_to_logical_devices(
    tensor, partition_dimensions
)

添加 tensor 将在逻辑设备之间拆分的注释。

这为张量 tensor 添加了一个注释，指定 tensor 上的操作将在多个逻辑设备之间拆分。张量 tensor 将在 partition_dimensions 指定的维度上拆分。 tensor 的尺寸必须能被 partition_dimensions 中的对应值整除。

例如，对于有 8 个逻辑设备的系统，如果 tensor 是一个具有形状 (batch_size, width, height, channel) 的图像张量，并且 partition_dimensions 是 [1, 2, 4, 1]，则 tensor将在宽度维度上拆分为 2，在高度维度上拆分为 4，拆分后的张量值将被馈送到 8 个逻辑设备中。

# Initializing TPU system with 8 logical devices and 1 replica.
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, 2, 2, 2],
    num_replicas=1)
# Construct the TPUStrategy. Since we are going to split the image across
# logical devices, here we set `experimental_spmd_xla_partitioning=True`
# so that the partitioning can be compiled in SPMD mode, which usually
# results in faster compilation and smaller HBM requirement if the size of
# input and activation tensors are much bigger than that of the model
# parameters. Note that this flag is suggested but not a hard requirement
# for `experimental_split_to_logical_devices`.
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment,
    experimental_spmd_xla_partitioning=True)

iterator = iter(inputs)

@tf.function()
def step_fn(inputs):
  inputs = strategy.experimental_split_to_logical_devices(
    inputs, [1, 2, 4, 1])

  # model() function will be executed on 8 logical devices with `inputs`
  # split 2 * 4  ways.
  output = model(inputs)
  return output

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

Args：张量：输入张量进行注释。 partition_dimensions：一个未嵌套的整数列表，其大小等于tensor 的等级，指定如何对tensor 进行分区。 partition_dimensions 中所有元素的乘积必须等于每个副本的逻辑设备总数。