Гайд · TNWS AI
Как пропустить NaN-gradient через Optax apply_if_finite
Optax 0.2.6: wrapper пропускает нечисловой update, увеличивает счётчики и сохраняет параметры на проблемном шаге.
Задача и применимость
Актуальный запрос: «как использовать apply_if_finite Optax». Защитить параметры от одного нечислового gradient. apply_if_finite оборачивает optimizer, проверяет NaN/Inf и при ошибке возвращает нулевые updates до превышения заданного числа последовательных ошибок. Это аварийный барьер, а не лечение причины нестабильности. В production дополнительно фиксируйте номер шага, global norm до сбоя и долю пропущенных updates. Частые пропуски делают обучение формально работающим, но фактически не продвигают параметры и требуют остановки.
Материал проверен на 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 и следуйте документации именно его версии.
Пошаговое решение
- Оберните SGD через apply_if_finite с пределом двух ошибок.
- Выполните корректный update и проверьте last_finite.
- Передайте gradient с NaN.
- Подтвердите нулевой update, last_finite=False и рост total_notfinite.
Готовый проверяемый шаблон
import jax.numpy as jnp
import optax
params = {"w": jnp.array([1., 1.])}
tx = optax.apply_if_finite(optax.sgd(0.1), max_consecutive_errors=2)
state = tx.init(params)
good, state = tx.update({"w": jnp.array([1., 2.])}, state, params)
bad, state = tx.update({"w": jnp.array([jnp.nan, 1.])}, state, params)
assert bool(state.last_finite) is False
assert int(state.total_notfinite) == 1
assert jnp.array_equal(bad["w"], jnp.zeros(2))
assert jnp.isfinite(good["w"]).all()
print(good, bad, state.total_notfinite)
Вход и ожидаемый результат
Корректный шаг создал update [-0.1,-0.2]; NaN-шаг вернул [0,0]. Состояние отметило last_finite=False и один найденный нечисловой update.
Как проверить результат
Логируйте total_notfinite и входные batch IDs без чувствительных данных. При повторных ошибках останавливайте обучение и исследуйте loss, данные, precision и learning rate. Нулевой update не означает успешный train step.
Перенос в рабочую модель
В рабочем проекте зафиксируйте контракт 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, мониторинг дрейфа, ограничения входов и человеческий контроль.
Типичные ошибки и чего не делать
- Не молча игнорировать счётчик NaN.
- Не считать wrapper заменой диагностики.
- Не применять параметры, если политика после max errors не проверена.
Не публикуйте только удачный вывод без версии среды, форм и негативного теста. Не настраивайте модель по test set и не загружайте непроверенное состояние поверх production-модели.
FAQ
Что происходит при NaN?
До лимита возвращается нулевой update.
Что считает total_notfinite?
Общее число обнаруженных неfinite updates.
Можно ли продолжить обучение?
Технически да после единичного сбоя, но причину нужно расследовать.
Официальные источники
Частые вопросы
Что происходит при NaN?
До лимита возвращается нулевой update.
Что считает total_notfinite?
Общее число обнаруженных неfinite updates.
Можно ли продолжить обучение?
Технически да после единичного сбоя, но причину нужно расследовать.
Читайте также
Как исключить 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 только на каждом втором вызове.
Комментарии
Пока тихо. Скажите первое слово