Лучшая производительность с tf.function с AutoGraph

Смотрите на TensorFlow.org Запустите в Google Colab Изучайте код на GitHub Скачайте ноутбук

TF 2.0 объединяет в себе простоту eager execution и мощь TF 1.0. В центре этого слияния находится tf.function, позволяющий преобразовывать подмножество синтаксиса Python в переносимые высокопроизводительные графы TensorFlow.

AutoGraph - крутая новая функция tf.function, которая позволяет писать код графа с использованием естественного синтаксиса Python. Список возможностей Python, которые можно использовать с AutoGraph, см. в Возможности и ограничения AutoGraph мкр. Для дополнительной информации о tf.function см. RFC TF 2.0: функции, а не сессии. Для более подробной информации об AutoGraph, см. tf.autograph.

Этот учебник познакомит вас с базовыми функциями tf.function и AutoGraph.

Setup

Импортируйте TensorFlow 2.0:

import numpy as np
import tensorflow as tf

Декоратор tf.function

Когда вы аннотируете функцию с помощьюtf.function, вы все равно можете вызывать ее как любую другую функцию. Но она будет скомпилирована в граф, что означает, что вы получаете преимущества более быстрого выполнения, запуска на GPU или TPU или экспорта в SavedModel.

@tf.function
def simple_nn_layer(x, y):
  return tf.nn.relu(tf.matmul(x, y))


x = tf.random.uniform((3, 3))
y = tf.random.uniform((3, 3))

simple_nn_layer(x, y)
<tf.Tensor: shape=(3, 3), dtype=float32, numpy=
array([[1.2318088 , 0.7174709 , 0.2651385 ],
       [1.1413519 , 0.61937577, 0.27351683],
       [1.5958037 , 1.0882345 , 0.28576517]], dtype=float32)>

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

simple_nn_layer
<tensorflow.python.eager.def_function.Function at 0x7f2590084240>

Если ваш код использует несколько функций, вам не нужно аннотировать их все - любые функции, вызываемые из аннотированной функции, также будут работать в режиме графа.

def linear_layer(x):
  return 2 * x + 1


@tf.function
def deep_net(x):
  return tf.nn.relu(linear_layer(x))


deep_net(tf.constant((1, 2, 3)))
<tf.Tensor: shape=(3,), dtype=int32, numpy=array([3, 5, 7], dtype=int32)>

Функции могут быть быстрее, чем eager код, для графов с большим количеством маленьких операций. Но для графов с несколькими дорогими операциями (например, свертки) вы можете не увидеть большого ускорения.

import timeit
conv_layer = tf.keras.layers.Conv2D(100, 3)

@tf.function
def conv_fn(image):
  return conv_layer(image)

image = tf.zeros([1, 200, 200, 100])
# warm up
conv_layer(image); conv_fn(image)
print("Eager conv:", timeit.timeit(lambda: conv_layer(image), number=10))
print("Function conv:", timeit.timeit(lambda: conv_fn(image), number=10))
print("Обратите внимание, что для операций свертки нет большой разницы в производительности")

Eager conv: 0.002389161000110107
Function conv: 0.00379357700012406
Обратите внимание, что для операций свертки нет большой разницы в производительности

lstm_cell = tf.keras.layers.LSTMCell(10)

@tf.function
def lstm_fn(input, state):
  return lstm_cell(input, state)

input = tf.zeros([10, 10])
state = [tf.zeros([10, 10])] * 2
# warm up
lstm_cell(input, state); lstm_fn(input, state)
print("eager lstm:", timeit.timeit(lambda: lstm_cell(input, state), number=10))
print("function lstm:", timeit.timeit(lambda: lstm_fn(input, state), number=10))

eager lstm: 0.004897988999800873
function lstm: 0.003677853999988656

Используйте порядок выполнения Python

При использовании зависимого от данных порядка выполнения внутриtf.function, вы можете использовать операторы порядка выполнения Python, и AutoGraph преобразует их в соответствующие операции TensorFlow. Например, операторы if будут преобразованы вtf.cond () если они зависят от Tensor-а.

В примере ниже x - Tensor но команда if работает как ожидалось:

@tf.function
def square_if_positive(x):
  if x > 0:
    x = x * x
  else:
    x = 0
  return x


print('square_if_positive(2) = {}'.format(square_if_positive(tf.constant(2))))
print('square_if_positive(-2) = {}'.format(square_if_positive(tf.constant(-2))))
square_if_positive(2) = 4
square_if_positive(-2) = 0

Примечание: предыдущий пример использует простые условные выражения со скалярными значениями. Batching обычно используется в коде для продакшна.

AutoGraph поддерживает обычные команды Python такие как while, for, if, break, continue и return, с поддержкой вложенности. Это значит что вы можете использовать Tensor выражения в условиях команд while и if, или итерировать по Tensor-у в цикле for.

@tf.function
def sum_even(items):
  s = 0
  for c in items:
    if c % 2 > 0:
      continue
    s += c
  return s


sum_even(tf.constant([10, 12, 15, 20]))
<tf.Tensor: shape=(), dtype=int32, numpy=42>

AutoGraph также предоставляет низкоуровневый API для опытных пользователей. Например, мы можем использовать его, чтобы взглянуть на сгенерированный код.

print(tf.autograph.to_code(sum_even.python_function))
def tf__sum_even(items):
    with ag__.FunctionScope('sum_even', 'fscope', ag__.ConversionOptions(recursive=True, user_requested=True, optional_features=(), internal_convert_user_code=True)) as fscope:
        do_return = False
        retval_ = ag__.UndefinedReturnValue()
        s = 0

        def get_state_2():
            return (s,)

        def set_state_2(vars_):
            nonlocal s
            (s,) = vars_

        def loop_body(itr):
            nonlocal s
            c = itr
            continue_ = False

            def get_state():
                return (continue_,)

            def set_state(vars_):
                nonlocal continue_
                (continue_,) = vars_

            def if_body():
                nonlocal continue_
                continue_ = True

            def else_body():
                nonlocal continue_
                pass
            ag__.if_stmt(((ag__.ld(c) % 2) > 0), if_body, else_body, get_state, set_state, ('continue_',), 1)

            def get_state_1():
                return (s,)

            def set_state_1(vars_):
                nonlocal s
                (s,) = vars_

            def if_body_1():
                nonlocal s
                s = ag__.ld(s)
                s += c

            def else_body_1():
                nonlocal s
                pass
            ag__.if_stmt(ag__.not_(continue_), if_body_1, else_body_1, get_state_1, set_state_1, ('s',), 1)
        c = ag__.Undefined('c')
        continue_ = ag__.Undefined('continue_')
        ag__.for_stmt(ag__.ld(items), None, loop_body, get_state_2, set_state_2, ('s',), {'iterate_names': 'c'})
        try:
            do_return = True
            retval_ = ag__.ld(s)
        except:
            do_return = False
            raise
        return fscope.ret(retval_, do_return)


Вот пример более сложного порядка выполнения:

@tf.function
def fizzbuzz(n):
  for i in tf.range(n):
    if i % 3 == 0:
      tf.print('Fizz')
    elif i % 5 == 0:
      tf.print('Buzz')
    else:
      tf.print(i)

fizzbuzz(tf.constant(15))
Fizz
1
2
Fizz
4
Buzz
Fizz
7
8
Fizz
Buzz
11
Fizz
13
14

Keras и AutoGraph

AutoGraph доступен по умолчанию в нединамических моделях Keras. Для получения дополнительной информации смотритеtf.keras.

class CustomModel(tf.keras.models.Model):

  @tf.function
  def call(self, input_data):
    if tf.reduce_mean(input_data) > 0:
      return input_data
    else:
      return input_data // 2


model = CustomModel()

model(tf.constant([-2, -4]))
<tf.Tensor: shape=(2,), dtype=int32, numpy=array([-1, -2], dtype=int32)>

Побочные эффекты

Как и в активном режиме, вы можете использовать операции с побочными эффектами, такие какtf.assign или tf.print, используемые обычно внутриtf.function, и он будет вставлять необходимые управляющие зависимости, чтобы обеспечить их выполнение в нужном порядке.

v = tf.Variable(5)

@tf.function
def find_next_odd():
  v.assign(v + 1)
  if v % 2 == 0:
    v.assign(v + 1)


find_next_odd()
v
<tf.Variable 'Variable:0' shape=() dtype=int32, numpy=7>

Отладка

tf.function и AutoGraph работают, генерируя код и копируя его в графы TensorFlow. Этот механизм пока не поддерживает пошаговые отладчики, такие как pdb. Однако вы можете вызвать tf.config.run_functions_eagerly(True) временно включить eager execution внутри tf.function и использовать ваш любимый отладчик:

@tf.function
def f(x):
  if x > 0:
    # Попробуйте установить точку остановки тут!
    # Example:
    #   import pdb
    #   pdb.set_trace()
    x = x + 1
  return x

tf.config.experimental_run_functions_eagerly(True)

# Сейчас вы можете установить точки остановки и запустить код в отладчике.
f(tf.constant(1))

tf.config.experimental_run_functions_eagerly(False)
WARNING:tensorflow:From <ipython-input-15-8ec88249052e>:11: experimental_run_functions_eagerly (from tensorflow.python.eager.def_function) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.config.run_functions_eagerly` instead of the experimental version.

Продвинутый пример: цикл обучения в графе

Предыдущий раздел показал, что AutoGraph можно использовать внутри слоев и моделей Keras. Модели Keras также могут быть использованы в коде AutoGraph.

Этот пример показывает, как обучить простую модель Keras в MNIST со всем процессом обучения - загрузкой пакетов, вычислением градиентов, обновлением параметров, вычислением точности валидации и повторением до сходимости - выполняемом в графе.

Скачайте данные

def prepare_mnist_features_and_labels(x, y):
  x = tf.cast(x, tf.float32) / 255.0
  y = tf.cast(y, tf.int64)
  return x, y

def mnist_dataset():
  (x, y), _ = tf.keras.datasets.mnist.load_data()
  ds = tf.data.Dataset.from_tensor_slices((x, y))
  ds = ds.map(prepare_mnist_features_and_labels)
  ds = ds.take(20000).shuffle(20000).batch(100)
  return ds

train_dataset = mnist_dataset()

Определите модель

model = tf.keras.Sequential((
    tf.keras.layers.Reshape(target_shape=(28 * 28,), input_shape=(28, 28)),
    tf.keras.layers.Dense(100, activation='relu'),
    tf.keras.layers.Dense(100, activation='relu'),
    tf.keras.layers.Dense(10)))
model.build()
optimizer = tf.keras.optimizers.Adam()

Определите цикл обучения

compute_loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

compute_accuracy = tf.keras.metrics.SparseCategoricalAccuracy()


def train_one_step(model, optimizer, x, y):
  with tf.GradientTape() as tape:
    logits = model(x)
    loss = compute_loss(y, logits)

  grads = tape.gradient(loss, model.trainable_variables)
  optimizer.apply_gradients(zip(grads, model.trainable_variables))

  compute_accuracy(y, logits)
  return loss


@tf.function
def train(model, optimizer):
  train_ds = mnist_dataset()
  step = 0
  loss = 0.0
  accuracy = 0.0
  for x, y in train_ds:
    step += 1
    loss = train_one_step(model, optimizer, x, y)
    if step % 10 == 0:
      tf.print('Step', step, ': loss', loss, '; accuracy', compute_accuracy.result())
  return step, loss, accuracy

step, loss, accuracy = train(model, optimizer)
print('Final step', step, ': loss', loss, '; accuracy', compute_accuracy.result())
Step 10 : loss 1.75328898 ; accuracy 0.406
Step 20 : loss 1.12783611 ; accuracy 0.5385
Step 30 : loss 0.731114447 ; accuracy 0.608666658
Step 40 : loss 0.638588369 ; accuracy 0.663
Step 50 : loss 0.486427307 ; accuracy 0.6982
Step 60 : loss 0.256945968 ; accuracy 0.7255
Step 70 : loss 0.536894202 ; accuracy 0.746142864
Step 80 : loss 0.351541221 ; accuracy 0.7655
Step 90 : loss 0.275927156 ; accuracy 0.781555533
Step 100 : loss 0.452086568 ; accuracy 0.7935
Step 110 : loss 0.279404551 ; accuracy 0.802636385
Step 120 : loss 0.517148793 ; accuracy 0.811083317
Step 130 : loss 0.299755067 ; accuracy 0.819230795
Step 140 : loss 0.369258791 ; accuracy 0.826071441
Step 150 : loss 0.303483129 ; accuracy 0.8306
Step 160 : loss 0.19394058 ; accuracy 0.836062491
Step 170 : loss 0.247444928 ; accuracy 0.840588212
Step 180 : loss 0.323490024 ; accuracy 0.844444454
Step 190 : loss 0.226716429 ; accuracy 0.848210514
Step 200 : loss 0.22766985 ; accuracy 0.8517
Final step tf.Tensor(200, shape=(), dtype=int32) : loss tf.Tensor(0.22766985, shape=(), dtype=float32) ; accuracy tf.Tensor(0.8517, shape=(), dtype=float32)

Батчинг (разбивка на пакеты)

В реальных приложениях разбивка на пакеты важна для производительности. Лучший код для преобразования в AutoGraph - это код, в котором порядок выполнения определяется на уровне batch. Если вы принимаете решения на уровне отдельного example, попробуйте использовать API пакетной обработки для сохранения производительности.

Например, если у вас следующий код в Python:

def square_if_positive(x):
  return [i ** 2 if i > 0 else i for i in x]


square_if_positive(range(-5, 5))
[-5, -4, -3, -2, -1, 0, 1, 4, 9, 16]

Возможно, у вас возникнет соблазн написать его в TensorFlow следующим образом (и это будет работать!):

@tf.function
def square_if_positive_naive(x):
  result = tf.TensorArray(tf.int32, size=x.shape[0])
  for i in tf.range(x.shape[0]):
    if x[i] > 0:
      result = result.write(i, x[i] ** 2)
    else:
      result = result.write(i, x[i])
  return result.stack()


square_if_positive_naive(tf.range(-5, 5))
<tf.Tensor: shape=(10,), dtype=int32, numpy=array([-5, -4, -3, -2, -1,  0,  1,  4,  9, 16], dtype=int32)>

Но в этом случае, оказывается вы можете написать следующее:

def square_if_positive_vectorized(x):
  return tf.where(x > 0, x ** 2, x)


square_if_positive_vectorized(tf.range(-5, 5))
<tf.Tensor: shape=(10,), dtype=int32, numpy=array([-5, -4, -3, -2, -1,  0,  1,  4,  9, 16], dtype=int32)>

Ре-трассировка

Ключевые моменты:

  • Соблюдайте осторожность при вызове функций с нетензорными аргументами или с аргументами, которые изменяют размерности.
  • Декорируйте функции уровня модуля и методы классов уровня модуля и избегайте декорирования локальных функций или методов.

tf.function может дать вам значительное ускорение относительно eager execution, цена этому более медленный первый запуск. Это потому что при первом запуске функция также трассируется в граф TensorFlow. Построение и оптимизация графа обычно намного медленнее его исполнения:

import timeit


@tf.function
def f(x, y):
  return tf.matmul(x, y)

print(
    "Первый вызов:",
    timeit.timeit(lambda: f(tf.ones((10, 10)), tf.ones((10, 10))), number=1))

print(
    "Второй вызов:",
    timeit.timeit(lambda: f(tf.ones((10, 10)), tf.ones((10, 10))), number=1))
Первый вызов: 0.02717175599991606
Второй вызов: 0.0008589649999066751

Вы можете легко определить, когда произошла трассировка функции, добавив оператор print в начало функции. Поскольку любой код Python выполняется только во время трассировки, вы увидите результат print, при трассировке функции:

@tf.function
def f():
  print('Tracing!')
  tf.print('Executing')

print('Первый вызов:')
f()

print('Второй вызов:')
f()
Первый вызов:
Tracing!
Executing
Второй вызов:
Executing

tf.function может также ре-трассировать при вызове других нетензорных аргументов:

@tf.function
def f(n):
  print(n, 'Трассировка!')
  tf.print(n, 'Выполненние')

f(1)
f(1)

f(2)
f(2)
1 Трассировка!
1 Выполненние
1 Выполненние
2 Трассировка!
2 Выполненние
2 Выполненние

Ре-трассировка может также произойти когда тензорные аргументы поменяли размерность, если вы не указали input_signature:

@tf.function
def f(x):
  print(x.shape, 'Трассировка!')
  tf.print(x, 'Выполненние')

f(tf.constant([1]))
f(tf.constant([2]))

f(tf.constant([1, 2]))
f(tf.constant([3, 4]))
(1,) Трассировка!
[1] Выполненние
[2] Выполненние
(2,) Трассировка!
[1 2] Выполненние
[3 4] Выполненние

В дополнение, tf.function always всегда создает новую функцию графа со своим наборов трассировок при каждом вызове:

def f():
  print('Трассировка!')
  tf.print('Выполнение')

tf.function(f)()
tf.function(f)()
Трассировка!
Выполнение
Трассировка!
Выполнение

Это может привести к неожиданному поведению при использовании декоратора @tf.function во вложенной функции:

def outer():
  @tf.function
  def f():
    print('Трассировка!')
    tf.print('Выполнение')
  f()

outer()
outer()
Трассировка!
Выполнение
Трассировка!
Выполнение