Repolode
投稿する
google

google/orbax

JAXモデルの重みやoptimizer stateなどの状態を保存・復元するための、チェックポイント作成・永続化ユーティリティ。

GitHub で見る公式サイト
スター
541
フォーク
109
未解決の Issue
452
最終コミット
6 時間前

これは何？

Orbaxは、JAXユーザー向けに共通のチェックポイント作成・永続化ユーティリティを提供するPythonライブラリです。モデルの重みやoptimizer stateなどのpytree状態を保存・復元できるため、学習コードで独自のシリアライズ処理を実装する必要がありません。

非同期チェックポイント作成、標準型とカスタム型、柔軟なストレージ形式に対応しています。APIは、さまざまなフレームワークのニーズに合わせてカスタマイズ、組み合わせができるよう設計されています。PyPIでは、チェックポイント用のorbax-checkpointとエクスポート用のorbax-exportが別パッケージとして公開されています。

Flax、MaxText、T5X、Tunixなど、複数のJAXフレームワークがチェックポイント作成にOrbaxを利用しています。

主な機能

  • 01シンプルなsave/load APIでJAXのpytree状態を保存・復元できます。
  • 02非同期チェックポイント作成に対応しています。
  • 03標準型とカスタム型を扱え、柔軟なストレージ形式を利用できます。
  • 04チェックポイント用とエクスポート用に、PyPIパッケージが分かれています。

ユースケース

  • JAXの学習中にモデルの重みやoptimizer stateを保存する。
  • 以前のチェックポイントから学習ジョブを再開する。
  • 組み合わせ可能なコンポーネントを使って、JAX上に独自のチェックポイント処理を構築する。

クイックスタート

  1. 1

    PyPIからチェックポイント用パッケージをインストールする

    pip install orbax-checkpoint
  2. 2

    READMEのクイックスタートコードを使ってJAXのpytreeを保存・復元する

使う前に知っておきたいこと

  • JAXが必要なため、JAXモデルをすでに扱っている場合にのみ役立ちます。

トピック

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

よくある質問

orbax とは？
JAXモデルの重みやoptimizer stateなどの状態を保存・復元するための、チェックポイント作成・永続化ユーティリティ。 Orbaxは、JAXユーザー向けに共通のチェックポイント作成・永続化ユーティリティを提供するPythonライブラリです。モデルの重みやoptimizer stateなどのpytree状態を保存・復元できるため、学習コードで独自のシリアライズ処理を実装する必要がありません。
orbax は無料で使えますか？ライセンスは？
orbax は Apache-2.0 ライセンスのオープンソースです。商用製品で使う前にライセンス条件を確認してください。
orbax は今もメンテナンスされていますか？
活発にメンテナンス中。orbax の最終コミットは 2026年11月6日 です。
orbax は何に使えますか？
JAXの学習中にモデルの重みやoptimizer stateを保存する。 以前のチェックポイントから学習ジョブを再開する。 組み合わせ可能なコンポーネントを使って、JAX上に独自のチェックポイント処理を構築する。
orbax のインストール方法は？
PyPIからチェックポイント用パッケージをインストールする: pip install orbax-checkpoint READMEのクイックスタートコードを使ってJAXのpytreeを保存・復元する
orbax はどこで入手できますか？
ソースコードとインストール手順は GitHub にあります：https://github.com/google/orbax

概要はプロジェクトの README をもとに AI が生成しています。GitHub の数値は毎日更新されます。

関連プロジェクト