Repolode
Предложить
google

google/orbax

Утилиты для сохранения и восстановления состояния моделей JAX, например весов и состояния оптимизатора.

Открыть на GitHubСайт
Звёзды
541
Форки
109
Открытые issues
452
Последний коммит
вчера

Что это?

Orbax — это библиотека Python с общими утилитами для создания чекпойнтов и сохранения состояния для пользователей JAX. Она сохраняет и восстанавливает состояние pytree, например веса модели и состояние оптимизатора, избавляя от необходимости реализовывать собственную логику сериализации в коде обучения.

Библиотека поддерживает асинхронное создание чекпойнтов, стандартные и пользовательские типы, а также гибкие форматы хранения. API можно настраивать и комбинировать под задачи разных фреймворков. Проект опубликован на PyPI в виде отдельных пакетов для создания чекпойнтов (orbax-checkpoint) и экспорта (orbax-export).

Orbax используют для создания чекпойнтов несколько фреймворков на базе JAX, включая Flax, MaxText, T5X и Tunix.

Ключевые возможности

  • 01Сохранение и восстановление состояния pytree в JAX с помощью простого API save/load.
  • 02Поддержка асинхронного создания чекпойнтов.
  • 03Работа со стандартными и пользовательскими типами и гибкими форматами хранения.
  • 04Отдельные пакеты на PyPI для создания чекпойнтов и экспорта.

Сценарии использования

  • Сохранение весов модели и состояния оптимизатора во время обучения на JAX.
  • Возобновление обучения с предыдущего чекпойнта.
  • Создание собственного процесса сохранения чекпойнтов на базе JAX из комбинируемых компонентов.

Быстрый старт

  1. 1

    Установите пакет для создания контрольных точек из PyPI

    pip install orbax-checkpoint
  2. 2

    Сохраните и восстановите pytree JAX с помощью кода быстрого старта из README

Что важно знать

  • Требуется JAX, поэтому это полезно только в том случае, если вы уже работаете с моделями JAX.

Темы

#jax#checkpointing#machine-learning#deep-learning#serialization#checkpoint#flax

Частые вопросы

Что такое orbax?
Утилиты для сохранения и восстановления состояния моделей JAX, например весов и состояния оптимизатора. Orbax — это библиотека Python с общими утилитами для создания чекпойнтов и сохранения состояния для пользователей JAX. Она сохраняет и восстанавливает состояние pytree, например веса модели и состояние оптимизатора, избавляя от необходимости реализовывать собственную логику сериализации в коде обучения.
orbax можно использовать бесплатно? Какая у него лицензия?
orbax распространяется с открытым исходным кодом под лицензией Apache-2.0. Перед использованием в коммерческих продуктах проверьте условия лицензии.
orbax всё ещё поддерживается?
Активно поддерживается. Последний коммит в orbax — 6 нояб. 2026 г..
Для чего можно использовать orbax?
Сохранение весов модели и состояния оптимизатора во время обучения на JAX. Возобновление обучения с предыдущего чекпойнта. Создание собственного процесса сохранения чекпойнтов на базе JAX из комбинируемых компонентов.
Как установить orbax?
Установите пакет для создания контрольных точек из PyPI: pip install orbax-checkpoint Сохраните и восстановите pytree JAX с помощью кода быстрого старта из README
Где скачать orbax?
Исходный код и инструкции по установке — на GitHub: https://github.com/google/orbax

Описание создано ИИ на основе README проекта. Статистика GitHub обновляется ежедневно.

Похожие проекты