Keras экономит память, повторяя вычисления: когда стоит включать rematerialization

Если обучение модели упирается в память, Keras предлагает обменять часть сохранённых промежуточных результатов на повторные вычисления. Механизм называется rematerialization, или gradient checkpointing. Это полезный вариант для проверки до покупки более ёмкого ускорителя, но экономию придётся сопоставить со временем обучения.

Откуда берётся экономия

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

Функция keras.remat(f) возвращает обёртку над вычислением f с соответствующим правилом градиента. Для собственной составной операции это позволяет обозначить участок, промежуточные состояния которого можно пересчитать. В документации есть пример такого пользовательского слоя; это иллюстрация механизма, а не измеренный выигрыш для любой модели.

Как ограничить область

Для работы со слоями есть keras.RematScope. Режим full охватывает поддерживаемые операции, activations выбирает активации, larger_than — слои с выходами выше заданного порога, list_of_layers — указанные имена слоёв. None отключает механизм. Контекст в примерах документации охватывает и создание слоёв, и вызов layer(input).

Начать можно с узкой области. Допустим, в вашей модели крупный блок уже назван encoder_block. Кандидат для проверки — RematScope(mode="list_of_layers", layer_names=["encoder_block"]). Это условное имя: замените его реальным, а не копируйте как название встроенного слоя. Затем выполните соответствующие вызовы в контексте по образцу документации.

Что сравнить перед выбором

Редакционный приёмочный сценарий прост: зафиксируйте модель, данные, размер батча, backend и оборудование. Сначала запишите пиковую память и время шага без пересчёта. Затем повторите измерение с выбранной областью. Сравнивайте одинаковый участок после прогрева, иначе подготовка вычислений может исказить вывод.

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

Выбирайте пересчёт, если он позволяет уложиться в память при приемлемой цене шага. Если память уже не ограничивает запуск, дополнительная работа может оказаться бесполезной. Сначала измерьте узкую область, затем решайте, нужно ли расширять её на модель.