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

用法

experimental_assign_to_logical_device(
    tensor, logical_device_id
)

添加将tensor 分配给逻辑设备的注释。

这将向 tensor 添加注释，指定将在逻辑核心设备 ID logical_device_id 上调用 tensor 上的操作。使用模型并行时，默认行为是所有操作都放置在zero-th 逻辑设备上。

# Initializing TPU system with 2 logical devices and 4 replicas.
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=4)
strategy = tf.distribute.TPUStrategy(
    resolver, experimental_device_assignment=device_assignment)
iterator = iter(inputs)

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

  # Add operation will be executed on logical device 0.
  output = strategy.experimental_assign_to_logical_device(output, 0)
  return output

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