TNWS AI — новости нейросетей и подбор моделей

Гайды · · 4 мин

Optax 0.2.6JAX 0.11.2PythonOptaxJAXнейросети

Как исключить bias из weight decay через Optax

Optax 0.2.6: маска add_decayed_weights применяет decay к kernel и оставляет bias без изменения.

Задача и применимость

Актуальный запрос: «как исключить bias из weight decay Optax add_decayed_weights». Применить add_decayed_weights только к kernel. Маска имеет ту же структуру PyTree: True включает transform, False исключает лист. При нулевых gradients, weight_decay 0,1, kernel 2 и bias 3 выход должен содержать добавку 0,2 для kernel и 0 для bias. Сам transform ещё не меняет знак и не умножает на learning rate. Для production отдельно зафиксируйте структуру PyTree, dtype, момент вызова относительно JIT и порядок соседних transformations. Добавьте негативный тест на несовместимую маску, неверный счётчик или отсутствующий batch axis — в зависимости от задачи. Численный unit test должен выполняться при каждом обновлении Optax.

Материал проверен на 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__. Не смешивайте разные способы накопления, маскирования и усреднения: похожие названия имеют разные контракты. При переносе сначала закрепите ожидаемую математику unit-тестом и следуйте документации установленной версии.

Пошаговое решение

  1. Создайте PyTree kernel/bias и нулевые gradients.
  2. Задайте boolean mask с True только для kernel.
  3. Вызовите add_decayed_weights и передайте params в update.
  4. Проверьте, что bias остался нулевым в updates.

Готовый проверяемый шаблон

import jax
import jax.numpy as jnp
import optax
params = {"kernel": jnp.array([2.0]), "bias": jnp.array([3.0])}
grads = jax.tree.map(jnp.zeros_like, params)
mask = {"kernel": True, "bias": False}
tx = optax.add_decayed_weights(0.1, mask=mask)
updates, _ = tx.update(grads, tx.init(params), params)
assert jnp.allclose(updates["kernel"], jnp.array([0.2]))
assert jnp.allclose(updates["bias"], jnp.array([0.0]))
print(updates)
Читать целиком
 

Гайды · · 4 мин

Optax 0.2.6JAX 0.11.2PythonOptaxJAXнейросети

Как использовать Huber loss в Optax для регрессии с выбросами

Optax 0.2.6: Huber loss квадратичен около нуля и линеен после delta; проверяем ошибки 0,5 и 2.

Задача и применимость

Актуальный запрос: «как использовать Huber loss Optax для выбросов». Проверить переход Huber loss от квадратичного режима к линейному. При target 0, delta=1 и prediction errors 0,5 и 2 ожидаются потери 0,125 и 1,5. Функция возвращает element-wise массив. Delta задаёт границу, но не выбирается автоматически по данным. До выбора delta исследуйте распределение абсолютных остатков на training split. Не подбирайте границу по test set. Если target стандартизирован, сохраняйте параметры обратного преобразования и показывайте MAE также в исходных единицах.

Материал проверен на 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__. Не смешивайте loss-функции с похожими названиями: они различаются ожидаемыми labels, формой результата и численной нормировкой. Сначала закрепите logits/targets/reduction контракт тестом и следуйте документации установленной версии.

Пошаговое решение

  1. Подготовьте predictions и broadcast-compatible targets.
  2. Выберите delta в единицах target после preprocessing.
  3. Вызовите huber_loss с keyword delta.
  4. Проверьте обе стороны границы и явно агрегируйте batch.

Готовый проверяемый шаблон

import jax.numpy as jnp
import optax
predictions = jnp.array([0.5, 2.0])
targets = jnp.zeros(2)
losses = optax.huber_loss(predictions, targets, delta=1.0)
assert jnp.allclose(losses, jnp.array([0.125, 1.5]))
assert losses[1] < optax.l2_loss(predictions, targets)[1]
print(losses)
Читать целиком
 

Гайды · · 4 мин

Optax 0.2.6JAX 0.11.2PythonOptaxJAXнейросети

Как использовать sigmoid focal loss в Optax

Optax 0.2.6: focal loss для logits 0 и labels 0/1 при gamma=2; проверка симметричного базового значения.

Задача и применимость

Актуальный запрос: «как использовать sigmoid focal loss Optax дисбаланс классов». Проверить focal loss на нейтральных logits. При logit 0, labels 0 и 1, gamma=2 и alpha=None обе потери равны приблизительно 0,1732868: обычная BCE ln(2) умножается на focal factor 0,25. Это element-wise функция по независимым бинарным классам. Gamma и alpha требуют validation. Для multilabel данных считайте частоту positives отдельно по каждому классу и проверяйте, не подавляет ли alpha полезные отрицательные примеры. Сравнение с BCE должно использовать те же logits, weights и reduction.

Материал проверен на 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__. Не смешивайте loss-функции с похожими названиями: они различаются ожидаемыми labels, формой результата и численной нормировкой. Сначала закрепите logits/targets/reduction контракт тестом и следуйте документации установленной версии.

Пошаговое решение

  1. Подготовьте logits и binary/soft labels одинаковой формы.
  2. Начните с gamma=2 и alpha=None для прозрачного baseline.
  3. Вычислите element-wise focal loss.
  4. Сравните вклад лёгких и сложных примеров, затем выберите reduction.

Готовый проверяемый шаблон

import jax.numpy as jnp
import optax
logits = jnp.array([0., 0.])
labels = jnp.array([0., 1.])
losses = optax.sigmoid_focal_loss(logits, labels, alpha=None, gamma=2.0)
expected = jnp.log(2.) * 0.25
assert jnp.allclose(losses, expected)
assert jnp.allclose(losses[0], losses[1])
print(losses)
Читать целиком
 

Гайды · · 4 мин

Optax 0.2.6JAX 0.11.2PythonOptaxJAXнейросети

Как ограничить gradient каждого примера через Optax

Optax 0.2.6: per_example_global_norm_clip отдельно ограничивает batch-градиенты и возвращает их сумму и число clipping-событий.

Задача и применимость

Актуальный запрос: «как ограничить норму градиента каждого примера Optax per_example_global_norm_clip». Ограничить L2-норму каждого примера до суммирования. Входные листья обязаны иметь batch в оси 0. Для gradients [3,4] с norm 5 и [0,6,0,8] с norm 1 при threshold 1 первый пример масштабируется до [0,6,0,8], второй остаётся прежним; сумма равна [1,2,1,6], clipped count равен 1. Для production отдельно зафиксируйте структуру PyTree, dtype, момент вызова относительно JIT и порядок соседних transformations. Добавьте негативный тест на несовместимую маску, неверный счётчик или отсутствующий batch axis — в зависимости от задачи. Численный unit test должен выполняться при каждом обновлении Optax.

Материал проверен на 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__. Не смешивайте разные способы накопления, маскирования и усреднения: похожие названия имеют разные контракты. При переносе сначала закрепите ожидаемую математику unit-тестом и следуйте документации установленной версии.

Пошаговое решение

  1. Подготовьте per-example gradients с batch axis 0.
  2. Вызовите функцию с l2_norm_clip=1.
  3. Прочитайте сумму clipped gradients и счётчик.
  4. При необходимости разделите сумму на batch size отдельно и явно.

Готовый проверяемый шаблон

import jax.numpy as jnp
import optax
grads = [jnp.array([[3.0, 4.0], [0.6, 0.8]])]
summed, clipped = optax.per_example_global_norm_clip(grads, l2_norm_clip=1.0)
assert jnp.allclose(summed[0], jnp.array([1.2, 1.6]))
assert int(clipped) == 1
mean = summed[0] / grads[0].shape[0]
assert jnp.allclose(mean, jnp.array([0.6, 0.8]))
print(summed, clipped, mean)
Читать целиком
 

Гайды · · 4 мин

Optax 0.2.6JAX 0.11.2PythonOptaxJAXнейросети

Как накапливать gradients через Optax MultiSteps

Optax 0.2.6: MultiSteps усредняет два microbatch-gradient и обновляет SGD только на каждом втором вызове.

Задача и применимость

Актуальный запрос: «как накапливать градиенты Optax MultiSteps». Имитировать больший batch с MultiSteps(optax.sgd(0.1), every_k_schedule=2). При use_grad_mean=True gradients 2 и 4 усредняются до 3: первый вызов возвращает ноль, второй update равен -0,3. Счётчики mini_step и gradient_step позволяют отличить microstep от фактического optimizer step. Для production отдельно зафиксируйте структуру PyTree, dtype, момент вызова относительно JIT и порядок соседних transformations. Добавьте негативный тест на несовместимую маску, неверный счётчик или отсутствующий batch axis — в зависимости от задачи. Численный unit test должен выполняться при каждом обновлении Optax.

Материал проверен на 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__. Не смешивайте разные способы накопления, маскирования и усреднения: похожие названия имеют разные контракты. При переносе сначала закрепите ожидаемую математику unit-тестом и следуйте документации установленной версии.

Пошаговое решение

  1. Оберните базовый SGD через MultiSteps с k=2.
  2. Передайте gradient 2 и убедитесь в нулевом update.
  3. Передайте gradient 4 и получите update среднего gradient.
  4. Проверьте счётчики и применяйте params только возвращённым update.

Готовый проверяемый шаблон

import jax.numpy as jnp
import optax
params = {"w": jnp.array([1.0])}
tx = optax.MultiSteps(optax.sgd(0.1), every_k_schedule=2)
state = tx.init(params)
u1, state = tx.update({"w": jnp.array([2.0])}, state, params)
u2, state = tx.update({"w": jnp.array([4.0])}, state, params)
assert jnp.allclose(u1["w"], 0.0)
assert jnp.allclose(u2["w"], -0.3)
assert int(state.mini_step) == 0 and int(state.gradient_step) == 1
print(u1, u2, state.mini_step, state.gradient_step)
Читать целиком
 

Гайды · · 4 мин

Optax 0.2.6JAX 0.11.2PythonOptaxJAXнейросети

Как обновлять target network через Optax periodic_update

Optax 0.2.6: periodic_update копирует online parameters на шагах, кратных update_period, и сохраняет target между ними.

Задача и применимость

Актуальный запрос: «как обновлять target network Optax periodic_update». Сделать жёсткое периодическое копирование online parameters в target network. При update_period=3 функция принимает new_tensors на шагах 0,3,6…; между ними возвращает old_tensors. Нулевая индексация означает, что копирование происходит уже на step=0. Для production отдельно зафиксируйте структуру PyTree, dtype, момент вызова относительно JIT и порядок соседних transformations. Добавьте негативный тест на несовместимую маску, неверный счётчик или отсутствующий batch axis — в зависимости от задачи. Численный unit test должен выполняться при каждом обновлении Optax.

Материал проверен на 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__. Не смешивайте разные способы накопления, маскирования и усреднения: похожие названия имеют разные контракты. При переносе сначала закрепите ожидаемую математику unit-тестом и следуйте документации установленной версии.

Пошаговое решение

  1. Создайте online и target PyTree.
  2. Передайте step=0 и подтвердите немедленное копирование.
  3. На шагах 1 и 2 подтвердите сохранение прежнего target.
  4. На шаге 3 проверьте следующее копирование.

Готовый проверяемый шаблон

import jax.numpy as jnp
import optax
target = {"w": jnp.array([1.0])}
r0 = optax.periodic_update({"w": jnp.array([10.0])}, target, steps=0, update_period=3)
r1 = optax.periodic_update({"w": jnp.array([11.0])}, r0, steps=1, update_period=3)
r3 = optax.periodic_update({"w": jnp.array([13.0])}, r1, steps=3, update_period=3)
assert jnp.allclose(r0["w"], 10.0)
assert jnp.allclose(r1["w"], 10.0)
assert jnp.allclose(r3["w"], 13.0)
print(r0, r1, r3)
Читать целиком
 
TNWS AI — новости нейросетей и подбор моделей