Гайд · TNWS AI
Как сделать warmup и cosine decay learning rate в Optax
Optax 0.2.6: линейный warmup до 0,1 за два шага и cosine decay до нуля; проверка контрольных точек.
Задача и применимость
Актуальный запрос: «как использовать warmup_cosine_decay_schedule Optax». Соединить линейный разогрев learning rate и последующий cosine decay одной официальной schedule. При init 0, peak 0,1, warmup_steps 2 и decay_steps 10 значение растёт 0→0,05→0,1, затем убывает и достигает нуля к общему шагу 10. Контроль промежуточного шага после peak защищает от ошибки, когда warmup работает, но decay случайно настроен на неверную продолжительность или отсчитывается в эпохах.
Материал проверен на Optax 0.2.6, JAX/jaxlib 0.11.2 и Python 3.12 на CPU. Используется актуальный API from flax import nnx, а не устаревший экспериментальный import и не Linen. Пример не загружает веса или внешние данные и решает одну изолированную задачу. Численный тест подтверждает контракт кода в указанной версии, но не качество нейросети на рабочем датасете.
Перед запуском установите совместимые версии пакетов в отдельное окружение и выведите flax.__version__, jax.__version__ и jaxlib.__version__. Не смешивайте примеры NNX с похожими по названию конструкциями Linen: у них различаются жизненный цикл Module, представление состояния, инициализация и способы применения transformations. При переносе кода сначала выберите один API и следуйте документации именно его версии.
Пошаговое решение
- Задайте init, peak, warmup_steps, decay_steps и end_value.
- Получите значения шагов 0,1,2,5,10.
- Проверьте линейный рост warmup и снижение после peak.
- Подключите функцию к optimizer и сохраняйте его step state.
Готовый проверяемый шаблон
import jax.numpy as jnp
import optax
schedule = optax.warmup_cosine_decay_schedule(
init_value=0.0, peak_value=0.1, warmup_steps=2,
decay_steps=10, end_value=0.0,
)
v = jnp.array([schedule(i) for i in [0, 1, 2, 5, 10]])
assert jnp.allclose(v[:3], jnp.array([0., 0.05, 0.1]), atol=1e-7)
assert v[3] < v[2]
assert jnp.allclose(v[4], 0., atol=1e-7)
print(v)
Вход и ожидаемый результат
Контрольные значения были [0,0.05,0.1,≈0.06913,0]. Peak достигнут на шаге 2, после него начался cosine decay.
Как проверить результат
Проверьте условие warmup_steps < decay_steps. При resume загружайте optimizer/schedule count, иначе warmup запустится повторно. Сопоставляйте step с фактическим числом updates после accumulation/skipping.
Перенос в рабочую модель
В рабочем проекте зафиксируйте контракт batch: порядок осей, dtype, диапазон признаков и семантику target. Инициализацию, dropout и генерацию данных разделяйте по именованным RNG streams. Обучаемые nnx.Param, статистическое состояние и optimizer state сохраняйте и проверяйте раздельно. Перед обновлением Flax/JAX прогоняйте небольшой обезличенный regression batch и сравнивайте формы, loss, конечность значений и допустимое численное отклонение.
Критерии приёмки
- Версии Flax, JAX и jaxlib записаны в отчёте.
- Формы входа, параметров и результата совпадают с контрактом.
- Все численные outputs и gradients конечны.
- Train/eval режим задан явно для стохастических и stateful слоёв.
- Результат проверен независимой формулой, состоянием или контрольным вызовом.
- Негативный случай неверной формы или режима включён в тесты проекта.
Регрессионный протокол
Сохраните минимальные входы: один объект, небольшой batch, нулевые значения, несимметричные размеры осей и заведомо несовместимую форму. Сравнивайте массивы через allclose с документированным допуском. Воспроизводимость требует одинаковых seeds и одинакового порядка запросов RNG. Производительность оценивайте после JIT-прогрева и block_until_ready, отдельно от compile time. Один удачный запуск не заменяет validation на независимых данных.
Диагностика расхождений
При расхождении сначала проверьте layout и preprocessing, затем RNG streams, train/eval mode, типы переменных State, forward, loss и optimizer update. Упростите модель до одного слоя, отключите JIT и стохастику, сравните промежуточные shapes. Возвращайте компоненты по одной. Не увеличивайте допуск и не заменяйте NaN нулями до установления причины.
Контроль после внедрения
В эксплуатации ограничьте допустимые shapes и размер batch, чтобы запрос не исчерпал память и не вызвал поток перекомпиляций. Не логируйте исходные пользовательские тензоры, ключи и содержимое обучающих примеров. Храните предыдущий проверенный checkpoint и схему State для отката. После восстановления проверяйте контрольный forward до включения модели в трафик.
Границы применимости
CPU-проверка не доказывает скорость или точную численную идентичность GPU/TPU. Случайная инициализация не является обучением, снижение train loss не гарантирует обобщение, а правильная форма не подтверждает корректность данных. Для важных решений нужны validation/test, мониторинг дрейфа, ограничения входов и человеческий контроль.
Типичные ошибки и чего не делать
- Не перезапускать schedule count при resume.
- Не задавать warmup длиннее всего decay интервала.
- Не считать batch number и optimizer step одинаковыми при accumulation.
Не публикуйте только удачный вывод без версии среды, форм и негативного теста. Не настраивайте модель по test set и не загружайте непроверенное состояние поверх production-модели.
FAQ
Когда достигается peak?
На границе warmup_steps.
Что происходит затем?
Cosine decay к end_value.
Почему это не две ручные функции?
Официальная schedule объединяет переход и единый счётчик шагов.
Официальные источники
Частые вопросы
Когда достигается peak?
На границе warmup_steps.
Что происходит затем?
Cosine decay к end_value.
Почему это не две ручные функции?
Официальная schedule объединяет переход и единый счётчик шагов.
Читайте также
Как исключить bias из weight decay через Optax
Optax 0.2.6: маска add_decayed_weights применяет decay к kernel и оставляет bias без изменения.
Как ограничить gradient каждого примера через Optax
Optax 0.2.6: per_example_global_norm_clip отдельно ограничивает batch-градиенты и возвращает их сумму и число clipping-событий.
Как накапливать gradients через Optax MultiSteps
Optax 0.2.6: MultiSteps усредняет два microbatch-gradient и обновляет SGD только на каждом втором вызове.
Комментарии
Пока тихо. Скажите первое слово