Python tf.compat.v1.tpu.outside_compilation用法及代码示例

在任何当前 TPU 复制范围之外构建计算的一部分。

用法

tf.compat.v1.tpu.outside_compilation(
    computation:Callable[..., Any],
    *args,
    **kwargs
) -> Any

tf.tpu.outside_compilation() 用于在 CPU 上运行 computation 中的操作，而不是在 TPU 上运行。例如，用户可以通过将这些操作显式放置在 CPU 上来运行 TPU 不支持的操作(例如 tf.summary.write())。下面使用外部编译会将操作放在 CPU 上的 computation_with_string_ops 中。

示例用法：

def computation_with_string_ops(x):
  # strings types are not supported on TPU's and below ops must
  # run on CPU instead.
  output = tf.strings.format('1{}', x)
  return tf.strings.to_number(output)

def tpu_computation():
  # Expected output is 11.
  output = tf.tpu.outside_compilation(computation_with_string_ops, 1)

外部编译应该在 TPUReplicateContext 内部调用。也就是说，tf.tpu.outside_compilation() 应该在传递给 tpu.split_compile_and_replicate() 的函数内部调用——这在传递给 TPUStrategy run() 的函数内部调用外部编译时暗示。如果在 TPUReplicateContext 之外调用，则这只会返回 computation 的结果，因此将是 no-op。请注意，外部编译与tf.distribute.experimental.TPUStrategy.merge_call() 不同，因为外部编译中的逻辑是为每个副本单独复制和执行的。另一方面，merge_call() 需要 merge_fn 来聚合来自不同副本的输入，并且只执行一次。

对于放置在 TPU 设备中的变量，包括在 TPUStrategy 范围内创建的变量，外部编译逻辑不得包含变量读/写。对于放置在主机上的变量，这是通过 TPUEstimator 创建的变量的情况，只有在 TPU 计算中的任何其他操作都没有访问该变量时才允许读取/写入变量。来自外部编译集群的变量读/写在 TPU 计算中不可见，反之亦然。因此，如果外部编译逻辑包含这样的主变量读/写操作，并且如果这些变量也被 TPU 计算访问，那么这可能会导致死锁。

在内部，tf.tpu.outside_compilation() 将外部编译属性添加到 computation 中的所有操作。在稍后的图形传递期间，这些具有外部编译属性的操作被提取出来并复制到host-side 图形中。此提取 host-side 图的输入通过一对 XlaSendToHost 和 XlaRecvFromHost 操作从 TPU 计算图发送到主机图。请注意，使用 tf.tpu.outside_compilation() 可能会导致 TPU 和 CPU 之间的张量传输，从而导致不小的性能影响。