flax
Version 0.10.2
Descripción general | Instalación rápida | ¿Cómo se ve el lino? | Documentación
Lanzado en 2024, Flax NNX es una nueva API de lino simplificada que está diseñada para facilitar la creación, inspeccionar, depurar y analizar redes neuronales en JAX. Logra esto agregando soporte de primera clase para la semántica de referencia de Python. Esto permite a los usuarios expresar sus modelos utilizando objetos de pitón regulares, permitiendo el intercambio de referencias y la mutabilidad.
Flax NNX evolucionó a partir de la API de lino de lino, que fue lanzado en 2020 por ingenieros e investigadores de Google Brain en estrecha colaboración con el equipo de JAX.
Puede obtener más información sobre FLax NNX en el sitio de documentación de lino dedicado. Asegúrese de ver:
NOTA: La documentación de Linen de Flax tiene su propio sitio.
La misión del equipo de lino es servir al creciente ecosistema de investigación de la red neuronal Jax, tanto dentro del alfabeto como con la comunidad en general, y explorar los casos de uso donde Jax brilla. Utilizamos GitHub para casi toda nuestra coordinación y planificación, así como donde discutimos los próximos cambios de diseño. Agradecemos los comentarios sobre cualquiera de nuestros hilos de discusión, emisión y solicitud de solicitudes.
Puede hacer solicitudes de funciones, hacernos saber en qué está trabajando, informar cuestiones, hacer preguntas en nuestro foro de discusión de Flax GitHub.
Esperamos mejorar el lino, pero no anticipamos cambios de ruptura significativos en la API central. Utilizamos entradas en ChangeLog y advertencias de deprecación cuando sea posible.
En caso de que desee comunicarse con nosotros directamente, estamos en [email protected].
Flax es una biblioteca y un ecosistema de red neuronal de alto rendimiento para Jax que está diseñado para la flexibilidad : pruebe nuevas formas de capacitación bifurcando un ejemplo y modificando el bucle de entrenamiento, sin agregar características a un marco.
Flax se está desarrollando en estrecha colaboración con el equipo de Jax y viene con todo lo que necesita para comenzar su investigación, incluyendo:
API de la red neuronal ( flax.nnx ): incluyendo Linear , Conv , BatchNorm , LayerNorm , GroupNorm , atención ( MultiHeadAttention ), LSTMCell , GRUCell , Dropout .
Utilidades y patrones : capacitación replicada, serialización y puntos de control, métricas, prioridad en el dispositivo.
Ejemplos educativos : MNIST, inferencia/muestreo con el modelo de lenguaje Gemma (Transformer), Transformer LM1B.
Flax usa Jax, así que consulte las instrucciones de instalación de JAX en CPU, GPU y TPUS.
Necesitará Python 3.8 o posterior. Instale lino desde pypi:
pip install flax
Para actualizar a la última versión de Flax, puede usar:
pip install --upgrade git+https://github.com/google/flax.git
Para instalar algunas dependencias adicionales (como matplotlib ) que son requeridas pero no incluidas por algunas dependencias, puede usar:
pip install " flax[all] " Proporcionamos tres ejemplos utilizando la API de lino: un perceptrón simple de múltiples capas, un CNN y un autoenvoder.
Para obtener más información sobre la abstracción Module , consulte nuestros documentos, nuestra amplia introducción a la abstracción del módulo. Para demostraciones concretas adicionales de las mejores prácticas, consulte nuestras guías y notas de desarrollador.
Ejemplo de un 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 )Ejemplo de un 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 xEjemplo de un 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 repositorio:
@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},
}
En la entrada de Bibtex anterior, los nombres están en orden alfabético, el número de versión está destinado a ser de Flax/Version.py, y el año corresponde a la versión de código abierto del proyecto.
Flax es un proyecto de código abierto mantenido por un equipo dedicado en Google Deepmind, pero no es un producto oficial de Google.
