jax
v0.4.35
Mulai cepat | Transformasi | Panduan pemasangan | Perpustakaan jaringan saraf | Ubah log | Dokumen referensi
JAX adalah pustaka Python untuk komputasi array berorientasi akselerator dan transformasi program, dirancang untuk komputasi numerik berkinerja tinggi dan pembelajaran mesin skala besar.
Dengan versi Autograd yang diperbarui, JAX dapat secara otomatis membedakan fungsi asli Python dan NumPy. Ia dapat berdiferensiasi melalui loop, cabang, rekursi, dan penutupan, dan dapat mengambil turunan dari turunan dari turunan. Ini mendukung diferensiasi mode terbalik (alias propagasi mundur) melalui diferensiasi grad dan mode maju, dan keduanya dapat disusun secara sewenang-wenang ke urutan apa pun.
Yang baru adalah JAX menggunakan XLA untuk mengkompilasi dan menjalankan program NumPy Anda pada GPU dan TPU. Kompilasi terjadi secara default, dengan panggilan perpustakaan dikompilasi dan dieksekusi tepat pada waktunya. Namun JAX juga memungkinkan Anda mengkompilasi fungsi Python Anda sendiri secara tepat waktu ke dalam kernel yang dioptimalkan XLA menggunakan API satu fungsi, jit . Kompilasi dan diferensiasi otomatis dapat disusun secara sewenang-wenang, sehingga Anda dapat mengekspresikan algoritme canggih dan mendapatkan performa maksimal tanpa meninggalkan Python. Anda bahkan dapat memprogram beberapa GPU atau inti TPU sekaligus menggunakan pmap , dan membedakan semuanya.
Gali lebih dalam, dan Anda akan melihat bahwa JAX benar-benar merupakan sistem yang dapat diperluas untuk transformasi fungsi yang dapat disusun. Baik grad maupun jit adalah contoh transformasi tersebut. Lainnya adalah vmap untuk vektorisasi otomatis dan pmap untuk pemrograman paralel beberapa program tunggal (SPMD) dengan banyak akselerator, dan masih banyak lagi yang akan datang.
Ini adalah proyek penelitian, bukan produk resmi Google. Harapkan bug dan ujung yang tajam. Tolong bantu dengan mencobanya, laporkan bug, dan beri tahu kami pendapat Anda!
import jax . numpy as jnp
from jax import grad , jit , vmap
def predict ( params , inputs ):
for W , b in params :
outputs = jnp . dot ( inputs , W ) + b
inputs = jnp . tanh ( outputs ) # inputs to the next layer
return outputs # no activation on last layer
def loss ( params , inputs , targets ):
preds = predict ( params , inputs )
return jnp . sum (( preds - targets ) ** 2 )
grad_loss = jit ( grad ( loss )) # compiled gradient evaluation function
perex_grads = jit ( vmap ( grad_loss , in_axes = ( None , 0 , 0 ))) # fast per-example gradsLangsung gunakan notebook di browser Anda, yang terhubung ke GPU Google Cloud. Berikut beberapa buku catatan pemula:
grad untuk diferensiasi, jit untuk kompilasi, dan vmap untuk vektorisasiJAX sekarang berjalan di Cloud TPU. Untuk mencoba pratinjau, lihat Cloud TPU Colabs.
Untuk mendalami JAX lebih dalam:
Pada intinya, JAX adalah sistem yang dapat diperluas untuk mengubah fungsi numerik. Berikut empat transformasi yang menjadi perhatian utama: grad , jit , vmap , dan pmap .
grad JAX memiliki API yang kurang lebih sama dengan Autograd. Fungsi yang paling populer adalah grad untuk gradien mode terbalik:
from jax import grad
import jax . numpy as jnp
def tanh ( x ): # Define a function
y = jnp . exp ( - 2.0 * x )
return ( 1.0 - y ) / ( 1.0 + y )
grad_tanh = grad ( tanh ) # Obtain its gradient function
print ( grad_tanh ( 1.0 )) # Evaluate it at x = 1.0
# prints 0.4199743 Anda dapat membedakan pesanan apa pun dengan grad .
print ( grad ( grad ( grad ( tanh )))( 1.0 ))
# prints 0.62162673 Untuk autodiff tingkat lanjut, Anda dapat menggunakan jax.vjp untuk produk vektor-Jacobian mode terbalik dan jax.jvp untuk produk vektor Jacobian mode maju. Keduanya dapat disusun secara sewenang-wenang satu sama lain, dan dengan transformasi JAX lainnya. Berikut salah satu cara menyusunnya untuk membuat fungsi yang secara efisien menghitung matriks Hessian penuh:
from jax import jit , jacfwd , jacrev
def hessian ( fun ):
return jit ( jacfwd ( jacrev ( fun )))Seperti halnya Autograd, Anda bebas menggunakan diferensiasi dengan struktur kontrol Python:
def abs_val ( x ):
if x > 0 :
return x
else :
return - x
abs_val_grad = grad ( abs_val )
print ( abs_val_grad ( 1.0 )) # prints 1.0
print ( abs_val_grad ( - 1.0 )) # prints -1.0 (abs_val is re-evaluated)Lihat dokumen referensi tentang diferensiasi otomatis dan JAX Autodiff Cookbook untuk informasi lebih lanjut.
jit Anda dapat menggunakan XLA untuk mengkompilasi fungsi Anda secara end-to-end dengan jit , digunakan sebagai dekorator @jit atau sebagai fungsi tingkat tinggi.
import jax . numpy as jnp
from jax import jit
def slow_f ( x ):
# Element-wise ops see a large benefit from fusion
return x * x + x * 2.0
x = jnp . ones (( 5000 , 5000 ))
fast_f = jit ( slow_f )
% timeit - n10 - r3 fast_f ( x ) # ~ 4.5 ms / loop on Titan X
% timeit - n10 - r3 slow_f ( x ) # ~ 14.5 ms / loop (also on GPU via JAX) Anda dapat menggabungkan jit dan grad serta transformasi JAX lainnya sesuka Anda.
Menggunakan jit memberikan batasan pada jenis aliran kontrol Python yang dapat digunakan fungsi; lihat tutorial Aliran Kontrol dan Operator Logis dengan JIT untuk informasi lebih lanjut.
vmap vmap adalah peta vektorisasi. Ia memiliki semantik yang familiar dalam memetakan suatu fungsi di sepanjang sumbu array, namun alih-alih menjaga perulangan di luar, ia mendorong perulangan ke bawah ke dalam operasi primitif suatu fungsi untuk kinerja yang lebih baik.
Menggunakan vmap dapat menyelamatkan Anda dari keharusan membawa-bawa dimensi batch dalam kode Anda. Misalnya, pertimbangkan fungsi prediksi jaringan neural sederhana yang tidak dikumpulkan ini:
def predict ( params , input_vec ):
assert input_vec . ndim == 1
activations = input_vec
for W , b in params :
outputs = jnp . dot ( W , activations ) + b # `activations` on the right-hand side!
activations = jnp . tanh ( outputs ) # inputs to the next layer
return outputs # no activation on last layer Kita sering kali menulis jnp.dot(activations, W) untuk memungkinkan dimensi batch di sisi kiri activations , namun kita telah menulis fungsi prediksi khusus ini untuk diterapkan hanya pada vektor masukan tunggal. Jika kita ingin menerapkan fungsi ini ke sekumpulan input sekaligus, secara semantik kita cukup menulis
from functools import partial
predictions = jnp . stack ( list ( map ( partial ( predict , params ), input_batch )))Namun mendorong satu contoh melalui jaringan pada satu waktu akan menjadi lambat! Lebih baik komputasinya divektorkan, sehingga di setiap lapisan kita melakukan perkalian matriks-matriks daripada perkalian matriks-vektor.
Fungsi vmap melakukan transformasi itu untuk kita. Artinya, jika kita menulis
from jax import vmap
predictions = vmap ( partial ( predict , params ))( input_batch )
# or, alternatively
predictions = vmap ( predict , in_axes = ( None , 0 ))( params , input_batch ) kemudian fungsi vmap akan mendorong loop luar ke dalam fungsi tersebut, dan mesin kita akan mengeksekusi perkalian matriks-matriks persis seperti jika kita melakukan pengelompokan dengan tangan.
Cukup mudah untuk mengelompokkan jaringan neural sederhana secara manual tanpa vmap , namun dalam kasus lain, vektorisasi manual mungkin tidak praktis atau tidak mungkin dilakukan. Ambil contoh masalah penghitungan gradien per contoh secara efisien: yaitu, untuk sekumpulan parameter tetap, kita ingin menghitung gradien fungsi kerugian yang dievaluasi secara terpisah pada setiap contoh dalam satu batch. Dengan vmap , caranya mudah:
per_example_gradients = vmap ( partial ( grad ( loss ), params ))( inputs , targets ) Tentu saja, vmap dapat disusun secara sewenang-wenang dengan jit , grad , dan transformasi JAX lainnya! Kami menggunakan vmap dengan diferensiasi otomatis mode maju dan mundur untuk perhitungan matriks Jacobian dan Hessian yang cepat di jax.jacfwd , jax.jacrev , dan jax.hessian .
pmap Untuk pemrograman paralel beberapa akselerator, seperti beberapa GPU, gunakan pmap . Dengan pmap Anda menulis program multi-data program tunggal (SPMD), termasuk operasi komunikasi kolektif paralel yang cepat. Menerapkan pmap berarti fungsi yang Anda tulis dikompilasi oleh XLA (mirip dengan jit ), kemudian direplikasi dan dijalankan secara paralel di seluruh perangkat.
Berikut ini contoh pada mesin 8-GPU:
from jax import random , pmap
import jax . numpy as jnp
# Create 8 random 5000 x 6000 matrices, one per GPU
keys = random . split ( random . key ( 0 ), 8 )
mats = pmap ( lambda key : random . normal ( key , ( 5000 , 6000 )))( keys )
# Run a local matmul on each device in parallel (no data transfer)
result = pmap ( lambda x : jnp . dot ( x , x . T ))( mats ) # result.shape is (8, 5000, 5000)
# Compute the mean on each device in parallel and print the result
print ( pmap ( jnp . mean )( result ))
# prints [1.1566595 1.1805978 ... 1.2321935 1.2015157]Selain mengekspresikan peta murni, Anda dapat menggunakan operasi komunikasi kolektif cepat antar perangkat:
from functools import partial
from jax import lax
@ partial ( pmap , axis_name = 'i' )
def normalize ( x ):
return x / lax . psum ( x , 'i' )
print ( normalize ( jnp . arange ( 4. )))
# prints [0. 0.16666667 0.33333334 0.5 ] Anda bahkan dapat menyarangkan fungsi pmap untuk pola komunikasi yang lebih canggih.
Semuanya tersusun, jadi Anda bebas membedakannya melalui penghitungan paralel:
from jax import grad
@ pmap
def f ( x ):
y = jnp . sin ( x )
@ pmap
def g ( z ):
return jnp . cos ( z ) * jnp . tan ( y . sum ()) * jnp . tanh ( x ). sum ()
return grad ( lambda w : jnp . sum ( g ( w )))( x )
print ( f ( x ))
# [[ 0. , -0.7170853 ],
# [-3.1085174 , -0.4824318 ],
# [10.366636 , 13.135289 ],
# [ 0.22163185, -0.52112055]]
print ( grad ( lambda x : jnp . sum ( f ( x )))( x ))
# [[ -3.2369726, -1.6356447],
# [ 4.7572474, 11.606951 ],
# [-98.524414 , 42.76499 ],
# [ -1.6007166, -1.2568436]] Saat mode mundur membedakan fungsi pmap (misalnya dengan grad ), proses komputasi ke belakang diparalelkan seperti halnya proses maju.
Lihat contoh SPMD Cookbook dan pengklasifikasi SPMD MNIST dari awal untuk mengetahui lebih lanjut.
Untuk survei yang lebih menyeluruh tentang gotcha saat ini, beserta contoh dan penjelasannya, kami sangat menyarankan untuk membaca Gotchas Notebook. Beberapa yang menonjol:
is tidak dipertahankan). Jika Anda menggunakan transformasi JAX pada fungsi Python yang tidak murni, Anda mungkin melihat kesalahan seperti Exception: Can't lift Traced... atau Exception: Different traces at same level .x[i] += y , tidak didukung, tetapi ada alternatif fungsional. Di bawah jit , alternatif fungsional tersebut akan menggunakan kembali buffer di tempatnya secara otomatis.jax.lax .float32 ) secara default, dan untuk mengaktifkan nilai presisi ganda (64-bit, misalnya float64 ) seseorang perlu menyetel variabel jax_enable_x64 saat startup (atau menyetel variabel lingkungan JAX_ENABLE_X64=True ) . Di TPU, JAX menggunakan nilai 32-bit secara default untuk semuanya kecuali variabel sementara internal dalam operasi 'seperti matmul', seperti jax.numpy.dot dan lax.conv . Operasi tersebut memiliki parameter precision yang dapat digunakan untuk memperkirakan operasi 32-bit melalui tiga lintasan bfloat16, dengan kemungkinan mengakibatkan waktu proses lebih lambat. Operasi non-matmul pada TPU lebih rendah dibandingkan implementasi yang sering kali menekankan kecepatan dibandingkan akurasi, sehingga dalam praktiknya penghitungan pada TPU akan kurang tepat dibandingkan penghitungan serupa pada backend lainnya.np.add(1, np.array([2], np.float32)).dtype adalah float64 bukan float32 .jit , membatasi cara Anda menggunakan aliran kontrol Python. Anda akan selalu mendapatkan kesalahan besar jika terjadi kesalahan. Anda mungkin harus menggunakan parameter jit static_argnums , primitif aliran kontrol terstruktur seperti lax.scan , atau cukup gunakan jit pada subfungsi yang lebih kecil. | Linuxx86_64 | Linux aarch64 | Mac x86_64 | Mac aarch64 | Windowsx86_64 | Windows WSL2 x86_64 | |
|---|---|---|---|---|---|---|
| CPU | Ya | Ya | Ya | Ya | Ya | Ya |
| GPU NVIDIA | Ya | Ya | TIDAK | tidak ada | TIDAK | eksperimental |
| Google TPU | Ya | tidak ada | tidak ada | tidak ada | tidak ada | tidak ada |
| GPU AMD | Ya | TIDAK | eksperimental | tidak ada | TIDAK | TIDAK |
| GPU Apple | tidak ada | TIDAK | tidak ada | eksperimental | tidak ada | tidak ada |
| GPU Intel | eksperimental | tidak ada | tidak ada | tidak ada | TIDAK | TIDAK |
| Platform | instruksi |
|---|---|
| CPU | pip install -U jax |
| GPU NVIDIA | pip install -U "jax[cuda12]" |
| Google TPU | pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html |
| GPU AMD (Linux) | Gunakan Docker, roda yang dibuat sebelumnya, atau buat dari sumber. |
| GPU Mac | Ikuti instruksi Apple. |
| GPU Intel | Ikuti instruksi Intel. |
Lihat dokumentasi untuk informasi tentang strategi instalasi alternatif. Ini termasuk kompilasi dari sumber, menginstal dengan Docker, menggunakan versi CUDA lain, conda build yang didukung komunitas, dan jawaban atas beberapa pertanyaan umum.
Beberapa grup riset Google di Google DeepMind dan Alphabet mengembangkan dan berbagi perpustakaan untuk melatih jaringan saraf di JAX. Jika Anda menginginkan perpustakaan berfitur lengkap untuk pelatihan jaringan saraf dengan contoh dan panduan cara, cobalah Flax dan situs dokumentasinya.
Lihat bagian Ekosistem JAX di situs dokumentasi JAX untuk daftar pustaka jaringan berbasis JAX, yang mencakup Optax untuk pemrosesan dan pengoptimalan gradien, chex untuk kode dan pengujian yang andal, dan Equinox untuk jaringan saraf. (Tonton pembicaraan Ekosistem JAX NeurIPS 2020 di DeepMind di sini untuk detail tambahan.)
Mengutip repositori ini:
@software{jax2018github,
author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang},
title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs},
url = {http://github.com/jax-ml/jax},
version = {0.3.13},
year = {2018},
}
Pada entri bibtex di atas, nama disusun berdasarkan abjad, nomor versi dimaksudkan dari jax/version.py, dan tahun sesuai dengan rilis sumber terbuka proyek.
Versi JAX yang baru lahir, yang hanya mendukung diferensiasi dan kompilasi otomatis ke XLA, dijelaskan dalam makalah yang muncul di SysML 2018. Saat ini kami sedang berupaya untuk meliput ide dan kemampuan JAX dalam makalah yang lebih komprehensif dan terkini.
Untuk detail tentang JAX API, lihat dokumentasi referensi.
Untuk memulai sebagai pengembang JAX, lihat dokumentasi pengembang.
