Repolode
提交项目
google

google/orbax

为 JAX 用户提供模型检查点保存与恢复的通用工具库。

在 GitHub 查看官网
Star
541
Fork
109
未关闭 Issue
452
最近提交
昨天

这是什么？

Orbax 是一个面向 JAX 用户的 Python 检查点与持久化工具库。它负责保存和恢复 pytree 形式的状态，例如模型权重和优化器状态，训练代码无需自行实现序列化逻辑。

该库支持异步检查点、标准与自定义数据类型以及灵活的存储格式，API 设计强调可定制与可组合性，以适应不同框架的需求。项目在 PyPI 上以独立包形式发布，其中 orbax-checkpoint 用于检查点，orbax-export 用于模型导出。

Flax、MaxText、T5X、Tunix 等 JAX 框架都在使用 Orbax 进行检查点管理。

核心功能

  • 01通过简单的 save/load 接口保存和恢复 JAX pytree 状态。
  • 02支持异步检查点写入。
  • 03支持标准类型与自定义类型，并提供灵活的存储格式。
  • 04检查点与导出功能分别以独立 PyPI 包发布。

适用场景

  • 在 JAX 训练过程中保存模型权重和优化器状态。
  • 从之前的检查点恢复训练任务。
  • 基于可组合组件为 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 用户提供模型检查点保存与恢复的通用工具库。 Orbax 是一个面向 JAX 用户的 Python 检查点与持久化工具库。它负责保存和恢复 pytree 形式的状态，例如模型权重和优化器状态，训练代码无需自行实现序列化逻辑。
orbax 可以免费使用吗？用的什么许可证？
orbax 以 Apache-2.0 许可证开源。用于商业产品前，请先确认许可证条款。
orbax 还在维护吗？
活跃维护中。orbax 最近一次提交是在 2026年11月6日。
orbax 适合用来做什么？
在 JAX 训练过程中保存模型权重和优化器状态。 从之前的检查点恢复训练任务。 基于可组合组件为 JAX 构建自定义检查点流程。
orbax 怎么安装和运行？
从 PyPI 安装检查点包: pip install orbax-checkpoint 使用 README 中的快速入门代码保存和恢复 JAX pytree
在哪里获取 orbax？
源代码和安装说明在 GitHub 上：https://github.com/google/orbax

概要由 AI 根据项目 README 生成；GitHub 数据每日更新。

相关项目