Понимание декоратора TensorFlow @tf.function

Введение

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

В этом кратком руководстве мы объясним, как tf.function() повышает производительность и ознакомьтесь с некоторыми рекомендациями.

Декораторы Python и tf.функция()

В Python декоратор — это функция, которая изменяет поведение других функций. Например, предположим, что вы вызываете следующую функцию в ячейке записной книжки:

import tensorflow as tf

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

def some_costly_computation(x):
    aux = tf.eye(100, dtype=tf.dtypes.float32)
    result = tf.zeros(100, dtype = tf.dtypes.float32)
    for i in range(1,100):
        aux = tf.matmul(x,aux)/i
        result = result + aux
    return result

%timeit some_costly_computation(x)

16.2 ms ± 103 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Однако, если мы передаем затратную функцию в tf.function():

quicker_computation = tf.function(some_costly_computation)
%timeit quicker_computation(x)

Мы получаем quicker_computation() – новая функция, которая работает намного быстрее предыдущей:

4.99 ms ± 139 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Итак, tf.function() модифицирует some_costly_computation() и выводит quicker_computation() функция. Декораторы также модифицируют функции, поэтому было естественно сделать tf.function() тоже декоратор.

Использование нотации декоратора аналогично вызову tf.function(function):

@tf.function
def quick_computation(x):
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

%timeit quick_computation(x)

5.09 ms ± 283 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Каким tf.function() Работа?

Как же мы можем заставить некоторые функции работать в 2-3 раза быстрее?

Код TensorFlow можно запускать в двух режимах: нетерпеливый режим и графический режим. Режим Eager — это стандартный интерактивный способ запуска кода: каждый раз, когда вы вызываете функцию, она выполняется.

Однако графический режим немного отличается. В графическом режиме перед выполнением функции TensorFlow создает граф вычислений, представляющий собой структуру данных, содержащую операции, необходимые для выполнения функции. Граф вычислений позволяет TensorFlow упростить вычисления и найти возможности для распараллеливания. График также изолирует функцию от вышележащего кода Python, что позволяет эффективно выполнять ее на многих различных устройствах.

Функция, украшенная @tf.function выполняется в два этапа:

1. На первом этапе TensorFlow выполняет код Python для функции и компилирует граф вычислений, задерживая выполнение любой операции TensorFlow.
2. После этого запускается расчетный граф.

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

@tf.function
def quick_computation(x):
  print('Only prints the first time!')
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

quick_computation(x)
quick_computation(x)

Это приводит к:

Only prints the first time!

Ассоциация print() выполняется только один раз на этапе трассировки, когда запускается обычный код Python. Следующие вызовы функции выполняют только операции TenforFlow из графа вычислений (операции TensorFlow).

Тем не менее, если мы используем tf.print() вместо:

@tf.function
def quick_computation_with_print(x):
  tf.print("Prints every time!")
  aux = tf.eye(100, dtype=tf.dtypes.float32)
  result = tf.zeros(100, dtype = tf.dtypes.float32)
  for i in range(1,100):
    aux = tf.matmul(x,aux)/i
    result = result + aux
  return result

quick_computation_with_print(x)
quick_computation_with_print(x)

Prints every time!
Prints every time!

TensorFlow включает в себя tf.print() в графе вычислений, так как это операция TensorFlow, а не обычная функция Python.

Лучшие практики с @tf.function

Написание кода с помощью операций TensorFlow

Как мы только что показали, некоторые части кода игнорируются графом вычислений. Это затрудняет прогнозирование поведения функции при кодировании с помощью «нормального» кода Python, как мы только что видели на примере print(). Лучше кодировать функцию с помощью операций TensorFlow, когда это применимо, чтобы избежать неожиданного поведения.

Например, for и while петли могут или не могут быть преобразованы в эквивалентный цикл TensorFlow. Поэтому цикл «for» лучше писать как векторизованную операцию, если это возможно. Это улучшит производительность вашего кода и обеспечит правильную трассировку вашей функции.

В качестве примера рассмотрим следующее:

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

@tf.function
def function_with_for(x):
    summ = float(0)
    for row in x:
      summ = summ + tf.reduce_mean(row)
    return summ

@tf.function
def vectorized_function(x):
  result = tf.reduce_mean(x, axis=0)
  return tf.reduce_sum(result)


print(function_with_for(x))
print(vectorized_function(x))

%timeit function_with_for(x)
%timeit vectorized_function(x)

tf.Tensor(0.672811, shape=(), dtype=float32)
tf.Tensor(0.67281103, shape=(), dtype=float32)
1.58 ms ± 177 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
440 µs ± 8.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Код с операциями TensorFlow значительно быстрее.

Избегайте ссылок на глобальные переменные

Рассмотрим следующий код:

x = tf.Variable(2, dtype=tf.dtypes.float32)
y = 2

@tf.function
def power(x):
  return tf.pow(x,y)

print(power(x))

y = 3

print(power(x))

tf.Tensor(4.0, shape=(), dtype=float32)
tf.Tensor(4.0, shape=(), dtype=float32)

Впервые украшенная функция power() был вызван, выходное значение было ожидаемым 4. Однако во второй раз функция проигнорировала, что значение y был изменен. Это происходит потому, что значение глобальных переменных Python замораживается для функции после трассировки.

Лучшим способом было бы использовать tf.Variable() для всех ваших переменных и передайте их в качестве аргументов вашей функции.

x = tf.Variable(2, dtype=tf.dtypes.float32)
y = tf.Variable(2, dtype = tf.dtypes.float32)

@tf.function
def power(x,y):
  return tf.pow(x,y)

print(power(x,y))

y.assign(3)

print(power(x,y))

tf.Tensor(4.0, shape=(), dtype=float32)
tf.Tensor(8.0, shape=(), dtype=float32)

Отладка [электронная почта защищена]_s

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

Некоторыми распространенными проблемами являются ошибки типа и ошибки формы. Ошибки типов возникают при несоответствии типов переменных, участвующих в операции:

x = tf.Variable(1, dtype = tf.dtypes.float32)
y = tf.Variable(1, dtype = tf.dtypes.int32)

z = tf.add(x,y)

InvalidArgumentError: cannot compute AddV2 as input #1(zero-based) was expected to be a float tensor but is a int32 tensor [Op:AddV2]

Ошибки типа легко закрадываются и могут быть легко исправлены путем приведения переменной к другому типу:

y = tf.cast(y, tf.dtypes.float32)
z = tf.add(x, y) 
tf.print(z) 

Ошибки формы возникают, когда ваши тензоры не имеют формы, необходимой для вашей операции:

x = tf.random.uniform(shape=[100, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)
y = tf.random.uniform(shape=[1, 100], minval=-1, maxval=1, dtype=tf.dtypes.float32)

z = tf.matmul(x,y)

InvalidArgumentError: Matrix size-incompatible: In[0]: [100,100], In[1]: [1,100] [Op:MatMul]

Одним из удобных инструментов для исправления обоих видов ошибок является интерактивный отладчик Python, который можно автоматически вызвать в блокноте Jupyter с помощью %pdb. Используя это, вы можете закодировать свою функцию и запустить ее в некоторых распространенных случаях использования. В случае ошибки открывается интерактивная подсказка. Это приглашение позволяет вам перемещаться вверх и вниз по слоям абстракции в вашем коде и проверять значения, типы и формы ваших переменных TensorFlow.

Заключение

Мы видели, как TensorFlow tf.function() делает вашу работу более эффективной, и как @tf.function декоратор применяет функцию к вашей собственной.

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