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. Ниже приведённое использование внешней компиляции разместит операции в 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)

Внешняя компиляция должна вызываться внутри 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, логика внешней компиляции не должна включать чтение/запись переменных. Для переменных, размещённых на хосте, что имеет место, когда переменные созданы с помощью TPUEstimator, чтение/запись переменных разрешено только в том случае, если к переменной не обращаются никакие другие операции в вычислении TPU. Чтение/запись переменных из кластера внешней компиляции не видны из вычисления TPU и наоборот. Следовательно, если логика внешней компиляции содержит такие операции чтения/записи переменных хоста, и если к переменным также обращается вычисление TPU, это может привести к тупику.

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