Гайд · TNWS AI
Как векторизовать инференс нейросети через jax.vmap
vmap для функции одного объекта: in_axes для параметров и входа, сравнение с циклом и пакетной матричной реализацией.
Задача и применимость
Актуальный запрос: «как использовать jax vmap для batch inference». Превратить функцию для одного объекта в пакетную без ручного Python-цикла. Параметры сети общие для всех строк, поэтому их ось задаётся None, а вход отображается по оси 0. Результат vmap должен совпасть с stack одиночных вызовов и с матричной batch-формой forward.
Материал применим к JAX 0.11.2 и jaxlib 0.11.2. JAX предоставляет массивы, автоматическое дифференцирование и программные преобразования; он не является готовой библиотекой слоёв, датасетов или полного тренировочного цикла. Пример намеренно мал и выполняется локально на доступном backend без загрузки весов и внешних данных. Численный критерий доказывает поведение конкретного кода, но не качество рабочей нейросети.
Пошаговое решение
- Опишите forward_one для вектора признаков.
- Создайте vmap с
in_axes=(None,0). - Передайте batch формы
(4,3). - Сравните с циклом и прямой batch-реализацией.
Готовый проверяемый шаблон
import jax
import jax.numpy as jnp
params={"W":jnp.arange(6,dtype=jnp.float32).reshape(3,2)/10,"b":jnp.array([.1,-.2])}
def forward_one(p,x): return jax.nn.relu(x@p["W"]+p["b"])
batch=jnp.array([[1.,0.,0.],[0.,1.,0.],[0.,0.,1.],[1.,1.,1.]])
batched=jax.vmap(forward_one,in_axes=(None,0))(params,batch)
loop=jnp.stack([forward_one(params,x) for x in batch])
direct=jax.nn.relu(batch@params["W"]+params["b"])
assert batched.shape==(4,2)
assert jnp.allclose(batched,loop) and jnp.allclose(batched,direct)
print(batched)
Вход и ожидаемый результат
Четыре объекта по три признака превращаются в массив (4,2). Результаты vmap, Python stack и прямой матричной формулы совпадают. Это функциональная проверка, а не доказательство ускорения на любом размере.
Как проверить результат
Явно задайте in_axes: None для общих params и 0 для batch. Сравните численно с эталоном на малом наборе. Измеряйте скорость только после jit и прогрева, если это соответствует рабочему пути.
Перенос на настоящую нейросеть
Для рабочей модели сначала зафиксируйте контракт параметров и 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 возвращает производную написанной функции, а не подтверждает правильность математической постановки. Для критичных решений нужны предметные метрики, мониторинг и человеческий контроль.
Типичные ошибки и чего не делать
- Не ставьте ось 0 для параметров, если они не имеют batch-копий.
- Не путайте vmap с распределением по устройствам.
- Не обещайте ускорение без benchmark.
Не публикуйте только удачный вывод без негативного теста и версий среды. Не используйте тестовую выборку для настройки модели. Не загружайте непроверенные параметры и не передавайте в диагностику персональные данные. Для важных решений предусмотрите безопасный отказ.
FAQ
Что делает vmap?
Векторизует функцию по указанным осям входов.
Почему params получают None?
Они общие для всех объектов batch.
Это multi-device API?
Нет, vmap добавляет векторизацию, а не распределение по устройствам.
Официальные источники
Частые вопросы
Что делает vmap?
Векторизует функцию по указанным осям входов.
Почему params получают None?
Они общие для всех объектов batch.
Это multi-device API?
Нет, vmap добавляет векторизацию, а не распределение по устройствам.
Читайте также
Как инициализировать двухслойную нейросеть в JAX
Двухслойный MLP без фреймворка: random.split, масштаб весов, ReLU, PyTree параметров, одиночный и пакетный входы.
Как найти 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.
Комментарии
Пока тихо. Скажите первое слово