tf.compat.v1.tpu.outside_compilation

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

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

tf.tpu.outside_compilation() используется для выполнения операций в computation на CPU, а не на TPU. Например, пользователи могут выполнять операции, которые не поддерживаются на TPU (например, tf.summary.write()), явно размещая эти операции на CPU. Приведённый ниже пример использования компиляции вне TPU размещает операции в computation_with_string_ops на CPU.

Пример использования:

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)

Компиляцию вне TPU следует вызывать внутри TPUReplicateContext. То есть tf.tpu.outside_compilation() следует вызывать внутри функции, передаваемой в tpu.split_compile_and_replicate() — это подразумевается, если компиляция вне TPU вызывается внутри функции, передаваемой в TPUStrategy run(). Если вызов выполнен вне TPUReplicateContext, то просто возвращается результат computation, и поэтому вызов ничего не делает. Обратите внимание, что компиляция вне TPU отличается от tf.distribute.experimental.TPUStrategy.merge_call(): логика компиляции вне TPU реплицируется и выполняется отдельно для каждой реплики. С другой стороны, для merge_call() требуется merge_fn, чтобы агрегировать входные данные разных реплик; она выполняется только один раз.

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

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