Почему тензорный поток потребляет так много оперативной памяти даже при обучении на ТПУ и как освободить память, выделенную после тренировки

Question

Я использую tf-1.15 и использую tf.TPUEstimator ( здесь ) для обучения модели. Модель работает отлично, за исключением того, что она использует около 3 ГБ ОЗУ. Это терпимо само по себе, но я зацикливаюсь на обучении более 10 таких моделей, и ОЗУ никогда не освобождается, поэтому оно складывается, потребляет более 30 ГБ ОЗУ ... Я уже пробовал

del estimator
gc.collect()
tf.reset_default_graph()

кажется, не работает.

Какие-либо предложения о том, как можно уменьшить использование памяти tenorflow или, по крайней мере, освободить их после завершения одного тренинга? Большое спасибо!

Answer 1

Nvm решение найдено. Кажется, что TF будет удерживать ресурсы, пока процесс не завершится (поправьте меня, если я ошибся). Итак, как я могу решить проблему, это обернуть сеанс / оценщик и запустить его внутри другого Процесса:

def core_run(args, kwargs):
    """ running the TF """
    ...
    return results

from multiprocessing import Pool
# preparing args and kwargs
with Pool(1) as p:
   results = p.apply(core, args, kwargs)

Таким образом, TF придется освободить ресурсы, поскольку процесс, в котором он находится, достигнет конца