flax
Version 0.10.2
Übersicht | Schnelle Installation | Wie sieht Flachs aus? | Dokumentation
Flax NNX wurde 2024 veröffentlicht und ist eine neue vereinfachte Flachs -API, mit der neuronale Netzwerke in JAX einfacher zu erstellen, zu inspizieren, zu debuggen und zu analysieren. Dies erreicht dies, indem es erstklassige Unterstützung für die Python -Referenzsemantik hinzufügt. Auf diese Weise können Benutzer ihre Modelle mit regelmäßigen Python -Objekten ausdrücken und die Referenzfreigabe und -veränderlichkeit ermöglichen.
Flachs NNX wurde aus der Flachs -Leinen -API entwickelt, die 2020 von Ingenieuren und Forschern von Google Brain in enger Zusammenarbeit mit dem JAX -Team veröffentlicht wurde.
Weitere Informationen zu Flachs NNX finden Sie auf der dedizierten Flachs -Dokumentationsseite. Stellen Sie sicher, dass Sie auschecken:
HINWEIS: Die Dokumentation von Flax Linen hat eine eigene Website.
Die Mission des Flachsteams ist es, dem wachsenden Jax -Ökosystem für das Neural Network Research - sowohl innerhalb von Alphabet als auch mit der breiteren Gemeinschaft zu dienen, und die Anwendungsfälle zu erforschen, in denen Jax strahlt. Wir verwenden GitHub für fast alle unsere Koordination und Planung sowie für die Besprechung bevorstehender Designänderungen. Wir begrüßen Feedback zu unserer Diskussions-, Ausgabe- und Anfrage -Threads.
Sie können Feature -Anfragen stellen, uns wissen, woran Sie arbeiten, Probleme melden, Fragen in unserem Diskussionsforum von Flax Github stellen.
Wir erwarten, Flachs zu verbessern, aber wir erwarten keine signifikanten Bruchänderungen der Kern -API. Wir verwenden ChangeLog -Einträge und Abschreibungswarnungen, wenn möglich.
Für den Fall, dass Sie uns direkt erreichen möchten, sind wir unter [email protected].
Flax ist eine leistungsstarke Bibliothek und ein Ökosystem für neuronale Netzwerke für JAX, das für Flexibilität ausgelegt ist: Probieren Sie neue Schulungsformen aus, indem Sie ein Beispiel geben und die Trainingsschleife ändern und keine Funktionen zu einem Framework hinzufügen.
Flax wird in enger Zusammenarbeit mit dem JAX -Team entwickelt und enthält alles, was Sie benötigen, um Ihre Recherchen zu beginnen, einschließlich:
Neurales Netzwerk -API ( flax.nnx ): einschließlich Linear , Conv , BatchNorm , LayerNorm , GroupNorm , Aufmerksamkeit ( MultiHeadAttention ), LSTMCell , GRUCell , Dropout .
Dienstprogramme und Muster : Replikiertes Training, Serialisierung und Überprüfungen, Metriken, Vorabstieg auf dem Gerät.
Bildungsbeispiele : MNIST, Inferenz/Stichproben mit dem Gemma -Sprachmodell (Transformator), Transformator LM1B.
Flax verwendet JAX. Schauen Sie sich also die Anweisungen zur JAX -Installation zu CPUs, GPUs und TPUs an.
Sie benötigen Python 3.8 oder höher. Flachs von PYPI einbauen:
pip install flax
Um auf die neueste Version von Flachs zu upgraden, können Sie verwenden:
pip install --upgrade git+https://github.com/google/flax.git
Um einige zusätzliche Abhängigkeiten (wie matplotlib ) zu installieren, die von einigen Abhängigkeiten nicht enthalten sind, können Sie verwenden:
pip install " flax[all] " Wir stellen drei Beispiele mit der Flachs-API an: ein einfaches Multi-Layer-Perzeptron, ein CNN und einen Auto-Coder.
Um mehr über die Module zu erfahren, lesen Sie unsere Dokumente, unser breites Intro zur Modulabstraktion. Weitere konkrete Demonstrationen Best Practices finden Sie in unseren Führern und Entwicklernotizen.
Beispiel eines 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 )Beispiel eines 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 xBeispiel eines AutoCodierers:
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 )So zitieren Sie dieses Repository:
@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},
}
Im obigen Bibtex-Eintrag befinden sich die Namen in alphabetischer Reihenfolge, die Versionsnummer soll die von Flax/Version.py sein und das Jahr der Open-Source-Version des Projekts entspricht.
Flax ist ein Open -Source -Projekt, das von einem engagierten Team bei Google DeepMind geführt wird, aber kein offizielles Google -Produkt.
