Гайд · TNWS AI
Как сохранить состояние Flax NNX в pure dict и восстановить параметры
Flax 0.12.10: to_pure_dict, replace_by_pure_dict, update и контроль точного forward после восстановления параметров.
Задача и применимость
Актуальный запрос: «как сохранить state Flax NNX в pure dict». Сохранить обучаемое состояние модели в обычный вложенный словарь и восстановить его после намеренного изменения весов. nnx.to_pure_dict извлекает массивы из State, а nnx.replace_by_pure_dict заменяет значения в совместимой структуре State; затем nnx.update переносит восстановленное состояние в живой модуль. Это полезный детерминированный слой проверки перед подключением Orbax или другого checkpoint-хранилища.
Материал проверен на Flax 0.12.10, 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 и следуйте документации именно его версии.
Пошаговое решение
- Создайте Linear и получите контрольный output на фиксированном входе.
- Извлеките State и преобразуйте его через nnx.to_pure_dict.
- Замените kernel нулями и подтвердите изменение результата.
- Создайте текущий State, примените replace_by_pure_dict, затем nnx.update и сравните output с эталоном.
Готовый проверяемый шаблон
import jax.numpy as jnp
from flax import nnx
model = nnx.Linear(2, 1, rngs=nnx.Rngs(0))
x = jnp.ones((1, 2))
expected = model(x)
saved = nnx.to_pure_dict(nnx.state(model))
model.kernel[...] = jnp.zeros_like(model.kernel[...])
assert not jnp.allclose(model(x), expected)
state = nnx.state(model)
nnx.replace_by_pure_dict(state, saved)
nnx.update(model, state)
actual = model(x)
assert jnp.allclose(actual, expected)
assert set(saved) == {"kernel", "bias"}
print(expected, actual, saved["kernel"].shape)
Вход и ожидаемый результат
У Linear 2→1 словарь содержал ключи kernel и bias; форма kernel была (2,1). После обнуления kernel результат изменился, а после замены State и nnx.update снова совпал с исходным в допуске allclose.
Как проверить результат
Проверяйте точное множество путей, shapes, dtype и контрольный forward. Pure dict сам по себе не хранит GraphDef архитектуры: при изменении структуры сначала выполните явную миграцию и проверьте её на копии. Для файлового checkpoint дополнительно используйте атомарную запись и метаданные версии.
Перенос в рабочую модель
В рабочем проекте зафиксируйте контракт 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, мониторинг дрейфа, ограничения входов и человеческий контроль.
Типичные ошибки и чего не делать
- Не считать pure dict полным описанием архитектуры.
- Не применять словарь от другой структуры без проверки путей и форм.
- Не перезаписывать единственный production checkpoint до контрольного восстановления.
Не публикуйте только удачный вывод без версии среды, форм и негативного теста. Не настраивайте модель по test set и не загружайте непроверенное состояние поверх production-модели.
FAQ
Сохраняет ли pure dict GraphDef?
Нет, он представляет значения State обычными контейнерами.
Зачем нужен nnx.update?
Он переносит восстановленный State в живой NNX-объект.
Можно ли хранить словарь на диске напрямую?
Нужен безопасный сериализатор/checkpointer и метаданные версии; документация показывает интеграцию с Orbax.
Официальные источники
Частые вопросы
Сохраняет ли pure dict GraphDef?
Нет, он представляет значения State обычными контейнерами.
Зачем нужен nnx.update?
Он переносит восстановленный State в живой NNX-объект.
Можно ли хранить словарь на диске напрямую?
Нужен безопасный сериализатор/checkpointer и метаданные версии; документация показывает интеграцию с Orbax.
Читайте также
Как накопить средний loss через Flax NNX metrics.Average
Flax 0.12.10: Average учитывает каждый элемент двух batch, compute возвращает общий mean, reset очищает состояние.
Как нормализовать каналы через Flax NNX GroupNorm
Flax 0.12.10: GroupNorm для NHWC, четыре канала и две группы, проверка статистик без зависимости от batch.
Как нормализовать каждый канал через Flax NNX InstanceNorm
Flax 0.12.10: InstanceNorm для NHWC нормирует пространственные значения каждого канала и каждого изображения отдельно.
Комментарии
Пока тихо. Скажите первое слово