flax
Version 0.10.2
Visão geral | Instalação rápida | Como é o linho? | Documentação
Lançado em 2024, o linho NNX é uma nova API de linho simplificada projetada para facilitar a criação, inspecionação, depuração e analisa redes neurais no JAX. Ele alcança isso adicionando suporte de primeira classe à semântica de referência do Python. Isso permite que os usuários expressem seus modelos usando objetos Python regulares, permitindo o compartilhamento de referência e a mutabilidade.
O linho NNX evoluiu da API de linho, lançado em 2020 por engenheiros e pesquisadores do Google Brain em estreita colaboração com a equipe JAX.
Você pode aprender mais sobre o linho NNX no site de documentação de linho dedicado. Certifique -se de conferir:
Nota: A documentação do linho de linho tem seu próprio site.
A missão da equipe de linho é servir o crescente ecossistema de pesquisa de rede neural JAX - tanto no alfabeto quanto na comunidade mais ampla, e explorar os casos de uso onde Jax brilha. Usamos o Github para quase toda a nossa coordenação e planejamento, bem como onde discutimos as próximas mudanças de design. Congratulamo -nos com feedback sobre qualquer um dos nossos tópicos de discussão, emissão e solicitação.
Você pode fazer solicitações de recursos, informe -nos no que está trabalhando, relatar problemas, fazer perguntas em nosso fórum de discussão do Github.
Esperamos melhorar o linho, mas não antecipamos mudanças significativas na API principal. Usamos entradas do Changelog e avisos de depreciação, quando possível.
Caso que você queira entrar em contato diretamente, estamos no [email protected].
O linho é uma biblioteca de redes neurais de alto desempenho e ecossistema para Jax, projetado para flexibilidade : tente novas formas de treinamento, comprando um exemplo e modificando o loop de treinamento, não adicionando recursos a uma estrutura.
O linho está sendo desenvolvido em estreita colaboração com a equipe JAX e vem com tudo o que você precisa para iniciar sua pesquisa, incluindo:
API da rede neural ( flax.nnx ): incluindo Linear , Conv , BatchNorm , LayerNorm , GroupNorm , atenção ( MultiHeadAttention ), LSTMCell , GRUCell , Dropout .
Utilitários e padrões : treinamento replicado, serialização e verificação, métricas, pré -busca no dispositivo.
Exemplos educacionais : mnist, inferência/amostragem com o modelo de linguagem gemma (transformador), transformador lm1b.
O linho usa o JAX, então confira as instruções de instalação do JAX em CPUs, GPUs e TPUs.
Você precisará de Python 3.8 ou posterior. Instale o linho do Pypi:
pip install flax
Para atualizar para a versão mais recente do linho, você pode usar:
pip install --upgrade git+https://github.com/google/flax.git
Para instalar algumas dependências adicionais (como matplotlib ) que são necessárias, mas não incluídas por algumas dependências, você pode usar:
pip install " flax[all] " Fornecemos três exemplos usando a API de linho: um simples perceptron de várias camadas, uma CNN e um codificador automático.
Para saber mais sobre a abstração Module , confira nossos documentos, nossa ampla introdução à abstração do módulo. Para demonstrações concretas adicionais de melhores práticas, consulte nossos guias e notas de desenvolvedor.
Exemplo de um MLP:
class MLP ( nnx . Module ):
def __init__ ( self , din : int , dmid : int , dout : int , * , rngs : nnx . Rngs ):
self . linear1 = Linear ( din , dmid , rngs = rngs )
self . dropout = nnx . Dropout ( rate = 0.1 , rngs = rngs )
self . bn = nnx . BatchNorm ( dmid , rngs = rngs )
self . linear2 = Linear ( dmid , dout , rngs = rngs )
def __call__ ( self , x : jax . Array ):
x = nnx . gelu ( self . dropout ( self . bn ( self . linear1 ( x ))))
return self . linear2 ( x )Exemplo de uma CNN:
class CNN ( nnx . Module ):
def __init__ ( self , * , rngs : nnx . Rngs ):
self . conv1 = nnx . Conv ( 1 , 32 , kernel_size = ( 3 , 3 ), rngs = rngs )
self . conv2 = nnx . Conv ( 32 , 64 , kernel_size = ( 3 , 3 ), rngs = rngs )
self . avg_pool = partial ( nnx . avg_pool , window_shape = ( 2 , 2 ), strides = ( 2 , 2 ))
self . linear1 = nnx . Linear ( 3136 , 256 , rngs = rngs )
self . linear2 = nnx . Linear ( 256 , 10 , rngs = rngs )
def __call__ ( self , x ):
x = self . avg_pool ( nnx . relu ( self . conv1 ( x )))
x = self . avg_pool ( nnx . relu ( self . conv2 ( x )))
x = x . reshape ( x . shape [ 0 ], - 1 ) # flatten
x = nnx . relu ( self . linear1 ( x ))
x = self . linear2 ( x )
return xExemplo de um autoencoder:
Encoder = lambda rngs : nnx . Linear ( 2 , 10 , rngs = rngs )
Decoder = lambda rngs : nnx . Linear ( 10 , 2 , rngs = rngs )
class AutoEncoder ( nnx . Module ):
def __init__ ( self , rngs ):
self . encoder = Encoder ( rngs )
self . decoder = Decoder ( rngs )
def __call__ ( self , x ) -> jax . Array :
return self . decoder ( self . encoder ( x ))
def encode ( self , x ) -> jax . Array :
return self . encoder ( x )Para citar este repositório:
@software{flax2020github,
author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},
title = {{F}lax: A neural network library and ecosystem for {JAX}},
url = {http://github.com/google/flax},
version = {0.10.2},
year = {2024},
}
Na entrada do BibTex acima, os nomes estão em ordem alfabética, o número da versão deve ser o do linho/versão.py, e o ano corresponde ao lançamento de código aberto do projeto.
O linho é um projeto de código aberto mantido por uma equipe dedicada no Google DeepMind, mas não é um produto oficial do Google.
