tf.compat.v1.tpu.outside_compilation

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

tf.compat.v1.tpu.outside_compilation(
    computation, *args, **kwargs
)

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, что приведет к нетривиальному влиянию на производительность.