OpenAI Градиентная контрольная точка с Tensorflow Eager Execution

Question

Я недавно перешел на Tensorflow Eager (в настоящее время работаю с TF 1.8.0) и мне это очень нравится. Однако теперь у меня есть довольно большая модель, которая не вписывается в мою память GPU (GTX 1080Ti, 12 ГБ VRAM) при работе с лентой градиента, которая необходима для вычисления градиентов в TF. Прямой проход (т.е. без использования ленты градиента) работает нормально.

Я думал об использовании Gradient Checkpointing из OpenAI в надежде, что это поможет. Однако простое использование его, как описано в их Git, похоже, не помогает в Eager Execution, т.е.

import tensorflow as tf
import tensorflow.contrib.eager as tfe
import memory_saving_gradients
tf.__dict__["gradients"] = memory_saving_gradients.gradients_memory
# using gradients_memory or gradients_speed does not change anything
# tf.__dict__["gradients"] = memory_saving_gradients.gradients_speed

[...]
with tfe.GradientTape() as g:
    output = run_large_model()
    loss = calculate_loss_on_output(output)
grads = g.gradient(full, model.variables)
optimizer.apply_gradients(zip(grads, model.variables))

не хватает памяти, независимо от использования градиентной контрольной точки или нет.

Я предполагаю, что лента градиента все еще хранит все переменные и необходимую информацию для обратного прохода, и контрольная точка градиента не имеет никакого эффекта, потому что TF в режиме Eager фактически не строит график (из того, что я понимаю - или, по крайней мере, это другой график).

Есть ли у вас какой-либо опыт или идея, как это можно решить или что мне нужно сделать, чтобы использовать контрольную точку градиента также в режиме TF Eager?

Answer 1

Код контрольной точки градиента из openai основан на переписывании графа, поэтому он не поддерживает активное выполнение.

Библиотека tenorflow.contrib.layers имеет декоратор Recompute_grad, который эквивалентен, но поддерживается как в графическом, так и в быстром исполнении.