tf.compat.v1.tpu.outside_compilation

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

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

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

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