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 中所有元素的乘積必須等於每個副本的邏輯設備總數。