Гайд · TNWS AI
Как инициализировать двухслойную нейросеть в JAX
Двухслойный MLP без фреймворка: random.split, масштаб весов, ReLU, PyTree параметров, одиночный и пакетный входы.
Задача и применимость
Актуальный запрос: «как создать простую нейросеть JAX два слоя». Собрать минимальный MLP непосредственно на jax.numpy, чтобы понимать формы параметров до использования более высокого фреймворка. Сеть принимает три признака, имеет четыре скрытых нейрона с ReLU и два выхода. Параметры хранятся в dict — стандартном pytree, совместимом с grad, jit и tree utilities.
Материал применим к JAX 0.11.2 и jaxlib 0.11.2. JAX предоставляет массивы, автоматическое дифференцирование и программные преобразования; он не является готовой библиотекой слоёв, датасетов или полного тренировочного цикла. Пример намеренно мал и выполняется локально на доступном backend без загрузки весов и внешних данных. Численный критерий доказывает поведение конкретного кода, но не качество рабочей нейросети.
Пошаговое решение
- Разделите PRNG key для двух матриц весов.
- Создайте W1
(3,4), b1(4,), W2(4,2), b2(2,). - Опишите forward через матричные операции и jax.nn.relu.
- Проверьте формы одиночного и пакетного результата.
Готовый проверяемый шаблон
import jax
import jax.numpy as jnp
from jax import random
key = random.key(7)
k1, k2 = random.split(key)
params = {
"W1": random.normal(k1, (3,4)) / jnp.sqrt(3.),
"b1": jnp.zeros((4,)),
"W2": random.normal(k2, (4,2)) / jnp.sqrt(4.),
"b2": jnp.zeros((2,)),
}
def forward(p, x):
hidden = jax.nn.relu(x @ p["W1"] + p["b1"])
return hidden @ p["W2"] + p["b2"]
one = forward(params, jnp.ones((3,)))
batch = forward(params, jnp.ones((5,3)))
assert one.shape == (2,) and batch.shape == (5,2)
assert all(jnp.isfinite(x).all() for x in jax.tree.leaves(params))
print(one, batch.shape)
Вход и ожидаемый результат
Один вектор из трёх признаков даёт два logits. Пакет из пяти строк формы (5,3) даёт (5,2). Все четыре листа параметров конечны и имеют заранее проверенные формы.
Как проверить результат
Перед обучением распечатайте tree structure, формы листьев и число параметров. Проверьте batch из одной строки и несколько строк. Конкретное деление на квадратный корень является учебной инициализацией, а не универсальным выбором для любой активации.
Перенос на настоящую нейросеть
Для рабочей модели сначала зафиксируйте контракт параметров и batch: pytree structure, формы, dtype, порядок признаков и семантику выходов. Генерацию случайностей оформите как явный поток keys. Разделите чистые функции forward, loss и update, чтобы их можно было независимо тестировать, дифференцировать и компилировать. Данные train, validation и test не должны смешиваться. Любое изменение формы, precision, backend или версии JAX проверяйте на сохранённом обезличенном regression-наборе до развёртывания.
Критерии приёмки
- Все формы и dtype совпадают с контрактом.
- Каждый PRNG subkey используется в одной логической случайной операции.
- Loss и gradients конечны на положительном тесте.
- Негативный вход вызывает ожидаемую ошибку, а не случайный output.
- Обычный и преобразованный через jit/vmap путь численно согласованы, когда это применимо.
- Устройство, версии JAX/jaxlib и численный допуск записаны в отчёте.
Регрессионный протокол
Храните несколько малых входов: стандартный batch, один объект, нули, граничные значения и заведомо невалидную форму. Сравнивайте outputs и loss с эталоном через allclose с явно выбранным допуском. После смены backend проверяйте допустимую численную погрешность, а не требуйте безусловной битовой идентичности. Производительность измеряйте после компиляционного прогрева и завершения асинхронной работы. Не выбирайте лучший единичный замер: сообщайте распределение задержки и условия машины.
Диагностика расхождений
Если результат изменился, локализуйте слой: данные и preprocessing, key-поток, shapes/dtype, forward, loss, gradients, update, затем jit и backend. Сначала выполните функцию без jit на минимальном входе, затем верните transformations по одной. Сравните tree structure и каждый leaf. Не увеличивайте допуск и не применяйте nan_to_num, пока не выяснено, не скрывает ли это неверную форму, переполнение или повтор ключа.
Контроль после внедрения
В эксплуатации ограничьте допустимые batch shapes и объём входа, чтобы запрос не исчерпал память и не создал бесконечный поток перекомпиляций. Собирайте агрегированные метрики latency, compile count, долю невалидных входов и численные ошибки без записи исходных чувствительных данных. Сохраняйте предыдущий проверенный артефакт и код параметров для отката. Автоматическая замена весов требует той же независимой проверки, что и обычный выпуск.
Границы применимости
Пример на CPU не доказывает свойства GPU или TPU, а маленький массив не предсказывает производительность большой модели. JIT может ускорить повторяющиеся вычисления, но стоимость компиляции, передача данных и новые формы способны ухудшить end-to-end время. Autodiff возвращает производную написанной функции, а не подтверждает правильность математической постановки. Для критичных решений нужны предметные метрики, мониторинг и человеческий контроль.
Типичные ошибки и чего не делать
- Не используйте один random subkey для обеих матриц.
- Не путайте число признаков и размер batch.
- Не применяйте ReLU к финальным logits перед cross-entropy без причины.
Не публикуйте только удачный вывод без негативного теста и версий среды. Не используйте тестовую выборку для настройки модели. Не загружайте непроверенные параметры и не передавайте в диагностику персональные данные. Для важных решений предусмотрите безопасный отказ.
FAQ
Зачем dict параметров?
Это pytree, который JAX умеет дифференцировать и преобразовывать.
Что выдаёт сеть?
Два необработанных выходных значения на строку.
Где batch-ось?
Первая ось входа (batch, features).
Официальные источники
Частые вопросы
Зачем dict параметров?
Это pytree, который JAX умеет дифференцировать и преобразовывать.
Что выдаёт сеть?
Два необработанных выходных значения на строку.
Где batch-ось?
Первая ось входа `(batch, features)`.
Читайте также
Как найти NaN в JAX-нейросети через debug_nans
Локальная диагностика NaN в eager и jit: контекст jax.debug_nans, FloatingPointError, безопасная функция и запрет маскировки проблемы nan_to_num.
Как обновить веса нейросети через value_and_grad в JAX
Один проверяемый шаг gradient descent для PyTree параметров: MSE, value_and_grad, tree.map, формы градиентов и снижение loss.
Как обновлять PyTree параметров через jax.tree.map
Проверяемые PyTrees: flatten, treedef, tree.map по двум структурам, unflatten и контроль несовместимого дерева.
Комментарии
Пока тихо. Скажите первое слово