Biblioteca modular e escalável para treinar modelos de machine learning em JAX/Flax, criada para acelerar a experimentação em pesquisa.
359+2 esta semana36PythonApache-2.0Bem mantidoAtualizado ontem
google/orbax
Utilitários de checkpointing e persistência para salvar e restaurar estados de modelos JAX, como pesos e estado do otimizador.
- Estrelas
- 541
- Forks
- 109
- Issues abertas
- 452
- Último commit
- há 23 horas
O que é?
Orbax é uma biblioteca Python que oferece utilitários comuns de checkpointing e persistência para quem usa JAX. Ela cuida de salvar e restaurar estados de pytree, como pesos do modelo e estado do otimizador, para que o código de treinamento não precise implementar sua própria lógica de serialização.
A biblioteca oferece suporte a checkpointing assíncrono, tipos padrão e personalizados e formatos de armazenamento flexíveis. Sua API foi projetada para ser personalizável e composta conforme as necessidades de diferentes frameworks. O projeto é publicado no PyPI em pacotes separados para checkpointing (orbax-checkpoint) e exportação (orbax-export).
Vários frameworks JAX usam Orbax para checkpointing, incluindo Flax, MaxText, T5X e Tunix.
Recursos principais
- 01Salva e restaura estados de pytree do JAX com uma API simples de save/load.
- 02Oferece suporte a checkpointing assíncrono.
- 03Lida com tipos padrão e personalizados usando formatos de armazenamento flexíveis.
- 04Disponibiliza pacotes separados no PyPI para checkpointing e exportação.
Casos de uso
- Salvar pesos do modelo e o estado do otimizador durante treinamentos com JAX.
- Retomar um trabalho de treinamento a partir de um checkpoint anterior.
- Criar um fluxo de checkpointing personalizado sobre o JAX com componentes combináveis.
Início rápido
- 1
Instale o pacote de checkpointing do PyPI
pip install orbax-checkpoint - 2
Salve e restaure uma pytree do JAX com o código de início rápido do README
Antes de usar
- Requer JAX, então só é útil se você já trabalha com modelos JAX.
Tópicos
Perguntas frequentes
- O que é orbax?
- Utilitários de checkpointing e persistência para salvar e restaurar estados de modelos JAX, como pesos e estado do otimizador. Orbax é uma biblioteca Python que oferece utilitários comuns de checkpointing e persistência para quem usa JAX. Ela cuida de salvar e restaurar estados de pytree, como pesos do modelo e estado do otimizador, para que o código de treinamento não precise implementar sua própria lógica de serialização.
- orbax é gratuito? Qual é a licença?
- orbax é open source sob a licença Apache-2.0. Confira os termos da licença antes de usá-lo em produtos comerciais.
- orbax ainda é mantido?
- Manutenção ativa. O último commit de orbax foi em 7 de nov. de 2026.
- Para que posso usar orbax?
- Salvar pesos do modelo e o estado do otimizador durante treinamentos com JAX. Retomar um trabalho de treinamento a partir de um checkpoint anterior. Criar um fluxo de checkpointing personalizado sobre o JAX com componentes combináveis.
- Como instalar orbax?
- Instale o pacote de checkpointing do PyPI: pip install orbax-checkpoint Salve e restaure uma pytree do JAX com o código de início rápido do README
- Onde consigo orbax?
- O código-fonte e as instruções de instalação estão no GitHub: https://github.com/google/orbax
Resumo gerado por IA a partir do README do projeto. As estatísticas do GitHub são atualizadas diariamente.
Projetos relacionados
Código em JAX/Flax e checkpoints pré-treinados para modelos de classificação de imagens Vision Transformer (ViT) e MLP-Mixer de artigos do Google Research.
12,8 mil+10 esta semana1,5 milJupyter NotebookApache-2.0Atualizado mês passado
Keras 3 é um framework de deep learning que roda com JAX, TensorFlow, PyTorch ou OpenVINO para inferência.
64,4 mil+8 esta semana19,8 milPythonApache-2.0Bem mantidoAtualizado ontem
Biblioteca Python que converte código de machine learning entre frameworks como PyTorch, TensorFlow, JAX e NumPy usando ivy.transpile.
14,2 mil5,4 milPythonOtherAtualizado há 17 dias
Uma biblioteca de redes neurais para JAX que prioriza flexibilidade, com a API Flax NNX para criar, inspecionar e depurar modelos.
7,4 mil+4 esta semana849Jupyter NotebookApache-2.0Atualizado há 24 horas
Solvers numéricos de equações diferenciais para JAX, com suporte a ODEs, SDEs e CDEs, autodiferenciação e execução em GPU.
IA e Machine Learning#differential-equations#jax#scientific-computing#neural-differential-equations#ode-solver
2,1 mil+1 esta semana202PythonApache-2.0Atualizado há 2 meses