Оптимизация графа TensorFlow с помощью Grappler

Обзор

TensorFlow использует для выполнения вычислений как граф, так и активное выполнение. tf.Graph содержит набор объектов tf.Operation (операций), которые представляют единицы вычислений, и объекты tf.Tensor , которые представляют единицы данных, которые передаются между операциями.

Grappler — это система оптимизации графов по умолчанию в среде выполнения TensorFlow. Grappler применяет оптимизацию в режиме графа (внутри tf.function ), чтобы повысить производительность ваших вычислений TensorFlow за счет упрощения графа и других оптимизаций высокого уровня, таких как встраивание тел функций для обеспечения межпроцедурной оптимизации. Оптимизация tf.Graph также снижает пиковое использование памяти устройства и улучшает использование оборудования за счет оптимизации сопоставления узлов графа с вычислительными ресурсами.

Используйте tf.config.optimizer.set_experimental_options() для более точного управления оптимизацией tf.Graph .

Доступные оптимизаторы графиков

Grappler выполняет оптимизацию графов с помощью драйвера верхнего уровня под названием MetaOptimizer . В TensorFlow доступны следующие оптимизаторы графов:

• Оптимизатор свертывания констант — статически выводит значение тензоров, когда это возможно, путем свертывания постоянных узлов в графе и материализует результат с использованием констант.
• Арифметический оптимизатор — упрощает арифметические операции, устраняя общие подвыражения и упрощая арифметические операторы.
• Оптимизатор макета — оптимизирует тензорные макеты для более эффективного выполнения операций, зависящих от формата данных, таких как свертки.
• Оптимизатор Remapper — переназначает подграфы на более эффективные реализации, заменяя часто встречающиеся подграфы оптимизированными объединенными монолитными ядрами.
• Оптимизатор памяти — анализирует график для проверки пикового использования памяти для каждой операции и вставляет операции копирования памяти ЦП-ГП для замены памяти ГП на ЦП, чтобы уменьшить пиковое использование памяти.
• Оптимизатор зависимостей — удаляет или переупорядочивает зависимости управления, чтобы сократить критический путь для шага модели или включить другие оптимизации. Также удаляются узлы, которые фактически не работают, такие как Identity.
• Оптимизатор сокращения — удаляет узлы, которые не влияют на вывод графика. Обычно он запускается первым, чтобы уменьшить размер графа и ускорить обработку в других проходах Grappler.
• Оптимизатор функций — оптимизирует библиотеку функций программы TensorFlow и встраивает тела функций для обеспечения других межпроцедурных оптимизаций.
• Оптимизатор формы — оптимизирует подграфы, которые работают с формой и информацией, связанной с формой.
• Автопараллельный оптимизатор — автоматически распараллеливает графики, разделяя их по пакетному измерению. Этот оптимизатор выключен по умолчанию.
• Оптимизатор циклов — оптимизирует поток управления графом, поднимая неизменяемые циклами подграфы из циклов и удаляя избыточные операции стека в циклах. Также оптимизирует циклы со статически известным количеством поездок и удаляет статически известные мертвые ветвления в условных операторах.
• Оптимизатор распределителя с заданной областью — вводит распределители с заданной областью для уменьшения перемещения данных и консолидации некоторых операций.
• Закрепить на хост-оптимизаторе — переносит небольшие операции на ЦП. Этот оптимизатор выключен по умолчанию.
• Автоматический оптимизатор смешанной точности — преобразует типы данных в float16, где это применимо, для повышения производительности. В настоящее время применяется только к графическим процессорам.
• Средство отладки отладки — удаляет узлы, связанные с операциями отладки, такие как tf.debugging.Assert , tf.debugging.check_numerics и tf.print из графа. Этот оптимизатор выключен по умолчанию.

Настраивать

import numpy as np
import timeit
import traceback
import contextlib


import tensorflow as tf

Создайте менеджер контекста, чтобы легко переключать состояния оптимизатора.

@contextlib.contextmanager
def options(options):
  old_opts = tf.config.optimizer.get_experimental_options()
  tf.config.optimizer.set_experimental_options(options)
  try:
    yield
  finally:
    tf.config.optimizer.set_experimental_options(old_opts)

Сравните производительность выполнения с Grappler и без него

TensorFlow 2 и более поздние версии по умолчанию выполняются быстро. Используйте tf.function , чтобы переключить выполнение по умолчанию в режим Graph. Grappler запускается автоматически в фоновом режиме, чтобы применить описанную выше оптимизацию графа и повысить производительность выполнения.

Оптимизатор постоянного сворачивания

В качестве предварительного примера рассмотрим функцию, которая выполняет операции над константами и возвращает результат.

def test_function_1():
  @tf.function
  def simple_function(input_arg):
    print('Tracing!')
    a = tf.constant(np.random.randn(2000,2000), dtype = tf.float32)
    c = a
    for n in range(50):
      c = c@a
    return tf.reduce_mean(c+input_arg)

  return simple_function

Отключите оптимизатор сворачивания констант и выполните функцию:

with options({'constant_folding': False}):
  print(tf.config.optimizer.get_experimental_options())
  simple_function = test_function_1()
  # Trace once
  x = tf.constant(2.2)
  simple_function(x)
  print("Vanilla execution:", timeit.timeit(lambda: simple_function(x), number = 1), "s")

{'constant_folding': False, 'disable_model_pruning': False, 'disable_meta_optimizer': False}
Tracing!
Vanilla execution: 0.0018392090000816097 s

Включите оптимизатор свертывания констант и снова выполните функцию, чтобы наблюдать ускорение выполнения функции.

with options({'constant_folding': True}):
  print(tf.config.optimizer.get_experimental_options())
  simple_function = test_function_1()
  # Trace once
  x = tf.constant(2.2)
  simple_function(x)
  print("Constant folded execution:", timeit.timeit(lambda: simple_function(x), number = 1), "s")

{'constant_folding': True, 'disable_model_pruning': False, 'disable_meta_optimizer': False}
Tracing!
Constant folded execution: 0.0006749789999958011 s

Оптимизатор отладки стриптизерши

Рассмотрим простую функцию, которая проверяет числовое значение своего входного аргумента и возвращает его.

def test_function_2():
  @tf.function
  def simple_func(input_arg):
    output = input_arg
    tf.debugging.check_numerics(output, "Bad!")
    return output
  return simple_func

Во-первых, выполните функцию с отключенным оптимизатором отладки.

test_func = test_function_2()
p1 = tf.constant(float('inf'))
try:
  test_func(p1)
except tf.errors.InvalidArgumentError as e:
  traceback.print_exc(limit=2)

2021-09-22 20:34:55.871238: E tensorflow/core/kernels/check_numerics_op.cc:292] abnormal_detected_host @0x7f4878e00100 = {0, 1} Bad!
Traceback (most recent call last):
  File "/tmp/ipykernel_22954/3616845043.py", line 4, in <module>
    test_func(p1)
  File "/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 885, in __call__
    result = self._call(*args, **kwds)
tensorflow.python.framework.errors_impl.InvalidArgumentError:  Bad! : Tensor had Inf values
     [[node CheckNumerics (defined at tmp/ipykernel_22954/2241890286.py:5) ]] [Op:__inference_simple_func_131]

Errors may have originated from an input operation.
Input Source operations connected to node CheckNumerics:
 input_arg (defined at tmp/ipykernel_22954/3616845043.py:4)

Function call stack:
simple_func

tf.debugging.check_numerics вызывает ошибку недопустимого аргумента из-за аргумента Inf для test_func .

Включите оптимизатор отладчика и снова выполните функцию.

with options({'debug_stripper': True}):
  test_func2 = test_function_2()
  p1 = tf.constant(float('inf'))
  try:
    test_func2(p1)
  except tf.errors.InvalidArgumentError as e:
    traceback.print_exc(limit=2)

Оптимизатор отладчика удаляет узел tf.debug.check_numerics из графа и выполняет функцию, не вызывая никаких ошибок.

Резюме

Среда выполнения TensorFlow использует Grappler для автоматической оптимизации графиков перед выполнением. Используйте tf.config.optimizer.set_experimental_options , чтобы включить или отключить различные оптимизаторы графиков.

Дополнительные сведения о Grappler см. в разделе Оптимизация графов TensorFlow .