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 之間的張量傳輸，從而導致不小的性能影響。