tf.compat.v1.tpu.outside_compilation

Создаёт часть вычисления вне текущего контекста репликации TPU.

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

tf.tpu.outside_compilation() используется для выполнения операций в computation на процессоре вместо выполнения на TPU. Например, пользователи могут выполнять операции, которые не поддерживаются на TPU (например, tf.summary.write()), явно размещая эти операции на процессорах. Ниже приведённое использование внешней компиляции разместит операции в 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, и поэтому будет являться пустой операцией. Обратите внимание, что внешняя компиляция отличается от tf.distribute.experimental.TPUStrategy.merge_call(), так как логика внешней компиляции дублируется и выполняется отдельно для каждой реплики. С другой стороны, merge_call() требует merge_fn для агрегирования входных данных из разных реплик и выполняется только один раз.

Для переменных, размещённых на устройстве TPU, что включает переменные, созданные внутри области TPUStrategy, логика внешней компиляции не должна включать чтение/запись переменных. Для переменных, размещённых на хосте, чтение/запись переменных разрешено только в том случае, если к переменной не обращаются никакие другие операции в вычислении TPU. Чтение/запись переменных из кластера внешней компиляции не видны из вычисления TPU и наоборот. Таким образом, если логика внешней компиляции содержит такие операции чтения/записи переменных хоста, и если к переменным также обращается вычисление TPU, это может привести к тупику.

Внутренне, tf.tpu.outside_compilation() добавляет атрибуты внешней компиляции ко всем операциям в computation. Во время последующих проходов операции с атрибутами внешней компиляции перемещаются в график на стороне хоста. Входные данные для этого графика на стороне хоста передаются из графика вычисления TPU на график хоста через пару операций XlaSendToHost и XlaRecvFromHost. Обратите внимание, что использование tf.tpu.outside_compilation() может привести к передаче тензоров между TPU и процессором, что приведёт к существенному влиянию на производительность.