tf.recompute_grad

Определяет функцию как контрольную точку пересчёта для автоматического дифференцирования с лентой.

tf.recompute_grad(
    f
)

Контрольные точки ленты — это метод снижения потребления памяти лентой автоматического дифференцирования:

• Без контрольных точек ленты операции и промежуточные значения записываются в ленту для использования при обратном проходе.

• При использовании контрольных точек ленты записываются только вызов функции и её входные данные. Во время обратного распространения пользовательский градиент recompute_grad (tf.custom_gradient) повторно вычисляет функцию в локальной области действия объекта Tape. Повторное вычисление функции при обратном распространении приводит к избыточным вычислениям, но снижает общее потребление памяти объектом Tape.

y = tf.Variable(1.0)

def my_function(x):
  tf.print('running')
  z = x*y
  return z

my_function_recompute = tf.recompute_grad(my_function)

with tf.GradientTape() as tape:
  r = tf.constant(1.0)
  for i in range(4):
    r = my_function_recompute(r)
running
running
running
running

grad = tape.gradient(r, [y])
running
running
running
running

Без recompute_grad лента содержит все промежуточные шаги, и повторное вычисление не выполняется.

with tf.GradientTape() as tape:
  r = tf.constant(1.0)
  for i in range(4):
    r = my_function(r)
running
running
running
running

grad = tape.gradient(r, [y])

Если f был объектом tf.keras Model или Layer, методы и атрибуты, такие как f.variables, недоступны в возвращённой функции g. Сохраните ссылку на f или используйте g.__wrapped__ для доступа к этим переменным и методам.

def print_running_and_return(x):
  tf.print("running")
  return x

model = tf.keras.Sequential([
  tf.keras.layers.Lambda(print_running_and_return),
  tf.keras.layers.Dense(2)
])

model_recompute = tf.recompute_grad(model)

with tf.GradientTape(persistent=True) as tape:
  r = tf.constant([[1,2]])
  for i in range(4):
    r = model_recompute(r)
running
running
running
running

grad = tape.gradient(r, model.variables)
running
running
running
running

В качестве альтернативы используйте атрибут __wrapped__, чтобы получить доступ к исходному объекту модели.

grad = tape.gradient(r, model_recompute.__wrapped__.variables)
running
running
running
running