Repolode
Enviar
google

google/flax

Uma biblioteca de redes neurais para JAX que prioriza flexibilidade, com a API Flax NNX para criar, inspecionar e depurar modelos.

Ver no GitHubSite
Estrelas
7,366
Forks
849
Issues abertas
526
Último commit
ontem

O que é?

Flax é uma biblioteca e um ecossistema de redes neurais criados sobre JAX. A API mais recente, Flax NNX, permite definir modelos com objetos Python comuns que oferecem compartilhamento de referências e mutabilidade, o que facilita criar, inspecionar, depurar e analisar redes. A API anterior, Flax Linen, lançada em 2020, é a predecessora da qual a NNX evoluiu.

A biblioteca oferece camadas comuns, como Linear, Conv, BatchNorm, LayerNorm, GroupNorm, atenção multi-head, células LSTM e GRU, além de Dropout. Também inclui utilitários para treinamento replicado, serialização e gerenciamento de checkpoints, métricas e pré-carregamento de dados nos dispositivos. Os exemplos abrangem MNIST e inferência com o modelo de linguagem Gemma.

Flax foi projetada para que os usuários modifiquem diretamente o loop de treinamento, em vez de depender de recursos adicionais do framework. A instalação está disponível pelo PyPI, e é necessário usar Python 3.8 ou posterior.

Recursos principais

  • 01API Flax NNX que usa objetos Python comuns com compartilhamento de referências e mutabilidade
  • 02Camadas integradas, incluindo Linear, Conv, camadas de normalização, atenção, células recorrentes e Dropout
  • 03Utilitários para treinamento replicado, gerenciamento de checkpoints, métricas e pré-carregamento de dados nos dispositivos
  • 04Exemplos de MNIST e inferência com o modelo de linguagem Gemma
  • 05Instalável pelo PyPI e baseado em JAX

Casos de uso

  • Pesquisar e criar protótipos de novas arquiteturas de redes neurais em JAX
  • Escrever loops de treinamento personalizados para experimentos sem restrições de framework
  • Treinar modelos em CPUs, GPUs ou TPUs usando JAX
  • Aprender deep learning com JAX usando os exemplos de MNIST e Gemma

Início rápido

  1. 1

    Primeiro, instale o JAX seguindo as instruções oficiais de instalação do JAX para o seu hardware

  2. 2

    Instale o Flax pelo PyPI

    pip install flax

Antes de usar

  • O Flax depende do JAX, então você também precisa instalar e configurar o JAX para o seu hardware.

Tópicos

#jax#neural-network#deep-learning#machine-learning#python#research

Perguntas frequentes

O que é flax?
Uma biblioteca de redes neurais para JAX que prioriza flexibilidade, com a API Flax NNX para criar, inspecionar e depurar modelos. Flax é uma biblioteca e um ecossistema de redes neurais criados sobre JAX. A API mais recente, Flax NNX, permite definir modelos com objetos Python comuns que oferecem compartilhamento de referências e mutabilidade, o que facilita criar, inspecionar, depurar e analisar redes. A API anterior, Flax Linen, lançada em 2020, é a predecessora da qual a NNX evoluiu.
flax é gratuito? Qual é a licença?
flax é open source sob a licença Apache-2.0. Confira os termos da licença antes de usá-lo em produtos comerciais.
flax ainda é mantido?
Manutenção ativa. O último commit de flax foi em 5 de nov. de 2026.
Para que posso usar flax?
Pesquisar e criar protótipos de novas arquiteturas de redes neurais em JAX Escrever loops de treinamento personalizados para experimentos sem restrições de framework Treinar modelos em CPUs, GPUs ou TPUs usando JAX Aprender deep learning com JAX usando os exemplos de MNIST e Gemma
Como instalar flax?
Primeiro, instale o JAX seguindo as instruções oficiais de instalação do JAX para o seu hardware Instale o Flax pelo PyPI: pip install flax
Onde consigo flax?
O código-fonte e as instruções de instalação estão no GitHub: https://github.com/google/flax

Resumo gerado por IA a partir do README do projeto. As estatísticas do GitHub são atualizadas diariamente.

Projetos relacionados