flax
Version 0.10.2
Aperçu | Installation rapide | À quoi ressemble le lin? | Documentation
Sorti en 2024, Flax NNX est une nouvelle API de lin simplifiée conçue pour faciliter la création, inspecter, déboguer et analyser les réseaux de neurones à Jax. Il y parvient en ajoutant une prise en charge de première classe pour la sémantique de référence Python. Cela permet aux utilisateurs d'exprimer leurs modèles à l'aide d'objets Python réguliers, permettant le partage de référence et la mutabilité.
Flax NNX a évolué à partir de l'API en lin de lin, qui a été publié en 2020 par des ingénieurs et des chercheurs de Google Brain en étroite collaboration avec l'équipe JAX.
Vous pouvez en savoir plus sur le lin NNX sur le site de documentation de lin dédié. Assurez-vous de vérifier:
Remarque: La documentation de Flax Linen a son propre site.
La mission de l'équipe de lin est de servir l'écosystème croissant de recherche sur le réseau neuronal de Jax - à la fois dans l'alphabet et avec la communauté plus large, et d'explorer les cas d'utilisation où Jax brille. Nous utilisons GitHub pour presque toute notre coordination et notre planification, ainsi que où nous discutons des changements de conception à venir. Nous accueillons des commentaires sur l'une de nos discussions sur les discussions, les problèmes et les demandes de traction.
Vous pouvez faire des demandes de fonctionnalités, faites-nous savoir sur quoi vous travaillez, signaler les problèmes, poser des questions dans notre forum de discussion sur le lin Github.
Nous nous attendons à améliorer le lin, mais nous ne prévoyons pas de changements de rupture significatifs à l'API principale. Nous utilisons les entrées de Changelog et les avertissements de dépréciation lorsque cela est possible.
Si vous souhaitez nous joindre directement, nous sommes à [email protected].
Le lin est une bibliothèque de réseaux neuronaux haute performance et un écosystème pour Jax conçu pour la flexibilité : essayez de nouvelles formes de formation en fournissant un exemple et en modifiant la boucle de formation, sans ajouter de fonctionnalités à un cadre.
Flax est en cours de développement en étroite collaboration avec l'équipe JAX et est livré avec tout ce dont vous avez besoin pour commencer vos recherches, notamment:
API du réseau neuronal ( flax.nnx ): y compris Linear , Conv , BatchNorm , LayerNorm , GroupNorm , Attention ( MultiHeadAttention ), LSTMCell , GRUCell , Dropout .
Utilitaires et modèles : formation répliquée, sérialisation et pointing de contrôle, mesures, préfecturation sur l'appareil.
Exemples éducatifs : MNIST, inférence / échantillonnage avec le modèle de langue Gemma (transformateur), Transformer LM1b.
Flax utilise JAX, alors consultez les instructions d'installation de JAX sur les CPU, les GPU et les TPU.
Vous aurez besoin de Python 3.8 ou version ultérieure. Installez le lin à partir de PYPI:
pip install flax
Pour passer à la dernière version de Flax, vous pouvez utiliser:
pip install --upgrade git+https://github.com/google/flax.git
Pour installer certaines dépendances supplémentaires (comme matplotlib ) qui sont nécessaires mais non incluses par certaines dépendances, vous pouvez utiliser:
pip install " flax[all] " Nous fournissons trois exemples en utilisant l'API de lin: un simple perceptron multicouche, un CNN et un encodeur automatique.
Pour en savoir plus sur l'abstraction Module , consultez nos documents, notre large introduction à l'abstraction du module. Pour des démonstrations concrets supplémentaires de meilleures pratiques, reportez-vous à nos guides et notes de développeur.
Exemple de 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 )Exemple de 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 xExemple d'un autoencodeur:
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 )Pour citer ce référentiel:
@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},
}
Dans l'entrée Bibtex ci-dessus, les noms sont dans l'ordre alphabétique, le numéro de version est censé être celui de Flax / Version.py, et l'année correspond à la version open source du projet.
Flax est un projet open source maintenu par une équipe dédiée à Google DeepMind, mais n'est pas un produit Google officiel.
