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