Гайд · TNWS AI
Как обновлять только kernel через Flax NNX PathContains
Flax 0.12.10: Optimizer с wrt=PathContains('kernel'), один SGD-шаг, kernel меняется, bias остаётся неизменным.
Задача и применимость
Актуальный запрос: «как заморозить bias Flax NNX PathContains». Заморозить bias и обновить только переменную, путь которой содержит kernel. Фильтр nnx.PathContains('kernel') передаётся в nnx.Optimizer через wrt; Optimizer использует только совпавшую часть градиентов. Это точная техника частичного fine-tuning, но имя пути должно быть проверено по реальной структуре модели.
Материал проверен на 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 и сохраните копии kernel и bias.
- Создайте Optimizer с wrt=nnx.PathContains('kernel').
- Вычислите полный gradient scalar loss.
- Выполните update и подтвердите изменение kernel при неизменном bias.
Готовый проверяемый шаблон
import jax.numpy as jnp
import optax
from flax import nnx
model = nnx.Linear(2, 1, rngs=nnx.Rngs(0))
kernel_before = model.kernel[...].copy()
bias_before = model.bias[...].copy()
optimizer = nnx.Optimizer(model, optax.sgd(0.1), wrt=nnx.PathContains("kernel"))
def loss_fn(m):
return jnp.mean(m(jnp.ones((3, 2))) ** 2)
grads = nnx.grad(loss_fn)(model)
optimizer.update(model, grads)
assert not jnp.array_equal(kernel_before, model.kernel[...])
assert jnp.array_equal(bias_before, model.bias[...])
print(model.kernel[...], model.bias[...])
Вход и ожидаемый результат
После одного SGD-шагa kernel изменился, а bias побитово совпал с копией до обновления. Тем самым проверена не только конфигурация фильтра, но и фактический эффект update.
Как проверить результат
До обучения выведите пути State и убедитесь, что фильтр выбирает ожидаемые листья и не пуст. После шага сравните замороженные и обучаемые массивы отдельно. В сложной модели узкое имя может совпасть с несколькими подмодулями — тестируйте полный набор путей.
Перенос в рабочую модель
В рабочем проекте зафиксируйте контракт 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, мониторинг дрейфа, ограничения входов и человеческий контроль.
Типичные ошибки и чего не делать
- Не считать requires_grad стилем другого фреймворка эквивалентом NNX filter.
- Не использовать слишком общее совпадение пути без аудита.
- Не забывать сохранить optimizer state именно для выбранного поддерева.
Не публикуйте только удачный вывод без версии среды, форм и негативного теста. Не настраивайте модель по test set и не загружайте непроверенное состояние поверх production-модели.
FAQ
Что выбирает PathContains?
Пути переменных, содержащие указанный ключ.
Считаются ли лишние gradients?
Документация отмечает, что при compiled step неиспользуемое вычисление может быть удалено.
Останется ли bias неизменным?
Да, если его путь не совпадает с фильтром и update выполнен этим Optimizer.
Официальные источники
Частые вопросы
Что выбирает PathContains?
Пути переменных, содержащие указанный ключ.
Считаются ли лишние gradients?
Документация отмечает, что при compiled step неиспользуемое вычисление может быть удалено.
Останется ли bias неизменным?
Да, если его путь не совпадает с фильтром и update выполнен этим Optimizer.
Читайте также
Как накопить средний 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 нормирует пространственные значения каждого канала и каждого изображения отдельно.
Комментарии
Пока тихо. Скажите первое слово