Unduh flax - Unduh Kode Sumber flax

flax

Kode sumber lainnya

Version 0.10.2

Unduh

Rami: Perpustakaan Jaringan Saraf dan Ekosistem untuk Jax Dirancang untuk Fleksibilitas

Membangun

Ikhtisar | Instalasi Cepat | Seperti apa rami itu? | Dokumentasi

Dirilis pada tahun 2024, Flax NNX adalah API Flax Sederhana yang dirancang untuk membuatnya lebih mudah untuk membuat, memeriksa, men -debug, dan menganalisis jaringan saraf di JAX. Ini mencapai ini dengan menambahkan dukungan kelas satu untuk semantik referensi Python. Ini memungkinkan pengguna untuk mengekspresikan model mereka menggunakan objek Python biasa, memungkinkan berbagi referensi dan mutabilitas.

Flax NNX berevolusi dari Flax Linen API, yang dirilis pada tahun 2020 oleh para insinyur dan peneliti di Google Brain dalam kolaborasi erat dengan tim JAX.

Anda dapat mempelajari lebih lanjut tentang Flax NNX di situs dokumentasi Flax khusus. Pastikan Anda memeriksa:

Dasar -dasar Rami NNX
Tutorial MNIST
Mengapa Flax Nnx
Evolusi dari linen rami ke rami nnx

Catatan: Dokumentasi Flax Linen memiliki situsnya sendiri.

Misi tim Flax adalah untuk melayani ekosistem penelitian Jax Neural Network - baik dalam alfabet maupun dengan komunitas yang lebih luas, dan untuk mengeksplorasi kasus penggunaan di mana Jax bersinar. Kami menggunakan GitHub untuk hampir semua koordinasi dan perencanaan kami, serta di mana kami membahas perubahan desain yang akan datang. Kami menyambut umpan balik tentang salah satu diskusi kami, mengeluarkan dan tarik utas permintaan.

Anda dapat membuat permintaan fitur, beri tahu kami apa yang sedang Anda kerjakan, melaporkan masalah, mengajukan pertanyaan di forum diskusi rami GitHub kami.

Kami berharap dapat meningkatkan rami, tetapi kami tidak mengantisipasi perubahan yang signifikan pada API inti. Kami menggunakan entri changelog dan peringatan penghentian bila memungkinkan.

Jika Anda ingin menghubungi kami secara langsung, kami berada di [email protected].

Ringkasan

Flax adalah perpustakaan dan ekosistem jaringan saraf berkinerja tinggi untuk JAX yang dirancang untuk fleksibilitas : coba bentuk pelatihan baru dengan memperpanjang contoh dan dengan memodifikasi loop pelatihan, tidak menambahkan fitur ke dalam kerangka kerja.

Flax sedang dikembangkan dalam kolaborasi erat dengan tim JAX dan hadir dengan semua yang Anda butuhkan untuk memulai penelitian Anda, termasuk:

API Jaringan Saraf ( flax.nnx ): termasuk Linear , Conv , BatchNorm , LayerNorm , GroupNorm , perhatian ( MultiHeadAttention ), LSTMCell , GRUCell , Dropout .
Utilitas dan Pola : Pelatihan yang direplikasi, serialisasi dan pos pemeriksaan, metrik, prefetching pada perangkat.
Contoh pendidikan : MNIST, inferensi/pengambilan sampel dengan model bahasa Gemma (transformator), transformator LM1B.

Instalasi cepat

Flax menggunakan JAX, jadi lihat instruksi instalasi JAX pada CPU, GPU, dan TPU.

Anda akan membutuhkan Python 3.8 atau lebih baru. Instal rami dari PYPI:

 pip install flax

Untuk meningkatkan ke versi terbaru dari Flax, Anda dapat menggunakan:

 pip install --upgrade git+https://github.com/google/flax.git

Untuk menginstal beberapa dependensi tambahan (seperti matplotlib ) yang diperlukan tetapi tidak termasuk oleh beberapa dependensi, Anda dapat menggunakan:

pip install " flax[all] "

Seperti apa rami itu?

Kami memberikan tiga contoh menggunakan API Flax: Perceptron multi-layer sederhana, CNN dan Auto-encoder.

Untuk mempelajari lebih lanjut tentang abstraksi Module , periksa dokumen kami, intro luas kami untuk abstraksi modul. Untuk demonstrasi konkret tambahan dari praktik terbaik, lihat panduan dan catatan pengembang kami.

Contoh 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 )

Contoh 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 x

Contoh 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 )

Mengutip rami

Untuk mengutip repositori ini:

 @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},
}

Dalam entri Bibtex di atas, nama dalam urutan abjad, nomor versi dimaksudkan dari Flax/Version.py, dan tahun ini sesuai dengan rilis sumber terbuka proyek.