Una biblioteca modular y escalable para entrenar modelos de machine learning con JAX/Flax, diseñada para iterar rápidamente en investigación.
358+2 esta semana36PythonApache-2.0Bien mantenidoActualizado hace 18 horas
google/orbax
Utilidades de checkpointing y persistencia para guardar y restaurar el estado de modelos JAX, como los pesos y el estado del optimizador.
- Estrellas
- 541
- Forks
- 109
- Issues abiertos
- 452
- Último commit
- hace 11 horas
¿Qué es?
Orbax es una biblioteca de Python que ofrece utilidades comunes de checkpointing y persistencia para usuarios de JAX. Permite guardar y restaurar estados de pytree, como los pesos del modelo y el estado del optimizador, para que el código de entrenamiento no tenga que implementar su propia lógica de serialización.
La biblioteca admite checkpointing asíncrono, tipos estándar y personalizados, y formatos de almacenamiento flexibles. Su API está diseñada para poder personalizarse y combinarse según las necesidades de cada framework. El proyecto se publica en PyPI como paquetes independientes para checkpointing (orbax-checkpoint) y exportación (orbax-export).
Varios frameworks de JAX usan Orbax para el checkpointing, entre ellos Flax, MaxText, T5X y Tunix.
Funciones clave
- 01Guardar y restaurar el estado de pytree de JAX mediante una API sencilla de guardado y carga.
- 02Admite checkpointing asíncrono.
- 03Gestiona tipos estándar y personalizados con formatos de almacenamiento flexibles.
- 04Paquetes independientes en PyPI para checkpointing y exportación.
Casos de uso
- Guardar los pesos del modelo y el estado del optimizador durante ejecuciones de entrenamiento con JAX.
- Reanudar un trabajo de entrenamiento a partir de un checkpoint anterior.
- Crear un flujo de trabajo de checkpointing personalizado sobre JAX con componentes combinables.
Inicio rápido
- 1
Instala el paquete de checkpointing desde PyPI
pip install orbax-checkpoint - 2
Guarda y restaura un pytree de JAX con el código de inicio rápido del README
Antes de usarlo
- Requiere JAX, por lo que solo resulta útil si ya trabajas con modelos de JAX.
Temas
Preguntas frecuentes
- ¿Qué es orbax?
- Utilidades de checkpointing y persistencia para guardar y restaurar el estado de modelos JAX, como los pesos y el estado del optimizador. Orbax es una biblioteca de Python que ofrece utilidades comunes de checkpointing y persistencia para usuarios de JAX. Permite guardar y restaurar estados de pytree, como los pesos del modelo y el estado del optimizador, para que el código de entrenamiento no tenga que implementar su propia lógica de serialización.
- ¿orbax es gratuito? ¿Qué licencia tiene?
- orbax es open source bajo la licencia Apache-2.0. Revisa los términos de la licencia antes de usarlo en productos comerciales.
- ¿orbax sigue mantenido?
- Mantenimiento activo. El último commit de orbax fue el 6 nov 2026.
- ¿Para qué puedo usar orbax?
- Guardar los pesos del modelo y el estado del optimizador durante ejecuciones de entrenamiento con JAX. Reanudar un trabajo de entrenamiento a partir de un checkpoint anterior. Crear un flujo de trabajo de checkpointing personalizado sobre JAX con componentes combinables.
- ¿Cómo se instala orbax?
- Instala el paquete de checkpointing desde PyPI: pip install orbax-checkpoint Guarda y restaura un pytree de JAX con el código de inicio rápido del README
- ¿Dónde consigo orbax?
- El código fuente y las instrucciones de instalación están en GitHub: https://github.com/google/orbax
Resumen generado por IA a partir del README del proyecto. Las estadísticas de GitHub se actualizan a diario.
Proyectos relacionados
Código JAX/Flax y checkpoints preentrenados para modelos de clasificación de imágenes Vision Transformer (ViT) y MLP-Mixer de artículos de Google Research.
12,8 mil+8 esta semana1,5 milJupyter NotebookApache-2.0Actualizado el mes pasado
Keras 3 es un framework de deep learning que funciona con JAX, TensorFlow, PyTorch u OpenVINO para inferencia.
64,4 mil+21 esta semana19,8 milPythonApache-2.0Bien mantenidoActualizado ayer
Una biblioteca de Python que convierte código de machine learning entre frameworks como PyTorch, TensorFlow, JAX y NumPy mediante ivy.transpile.
14,2 mil5,4 milPythonOtherActualizado hace 16 días
Una biblioteca de redes neuronales para JAX centrada en la flexibilidad, con la API Flax NNX para crear, inspeccionar y depurar modelos.
7,4 mil+4 esta semana849Jupyter NotebookApache-2.0Actualizado ayer
Solvers numéricos de ecuaciones diferenciales para JAX, con soporte para ODE, SDE y CDE, autodiferenciación y ejecución en GPU.
IA y Machine Learning#differential-equations#jax#scientific-computing#neural-differential-equations#ode-solver
2,1 mil+2 esta semana202PythonApache-2.0Actualizado hace 2 meses