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),))