Google JAX nedir? Hızlandırıcılarda NumPy

Popüler açık kaynaklı TensorFlow makine öğrenimi platformunu güçlendiren yenilikler arasında otomatik farklılaştırma (otograd) ve XLA (Hızlandırılmış Doğrusal Cebir) derin öğrenme için derleyiciyi optimize ediyor.

Bu iki teknolojiyi bir araya getiren bir diğer proje olan Google JAX, hız ve performans açısından önemli avantajlar sunuyor. GPU’larda veya TPU’larda çalıştırıldığında JAX, çağrı yapan diğer programların yerini alabilir. Dizi, ancak programları çok daha hızlı çalışır. Ek olarak, sinir ağları için JAX kullanmak, TensorFlow gibi daha büyük bir çerçeveyi genişletmekten çok yeni işlevler eklemeyi çok daha kolay hale getirebilir.

Bu makale, faydalarına ve sınırlamalarına genel bir bakış, kurulum talimatları ve Colab’daki Google JAX hızlı başlangıcına ilk bakış dahil olmak üzere Google JAX’ı tanıtmaktadır.

Otograd nedir?

Autograd, Ryan Adams’ın Harvard Intelligent Probabilistic Systems Group’ta bir araştırma projesi olarak başlayan otomatik bir farklılaşma motorudur. Bu yazı itibariyle, motorun bakımı yapılıyor ancak artık aktif olarak geliştirilmiyor. Bunun yerine geliştiricileri, Autograd’ı XLA JIT derlemesi gibi ek özelliklerle birleştiren Google JAX üzerinde çalışıyor. bu otograd motor yerel Python ve NumPy kodunu otomatik olarak ayırt edebilir. Birincil amaçlanan uygulaması, gradyan tabanlı optimizasyondur.

TensorFlow’lar tf.GradientTape API, Autograd’a benzer fikirlere dayanmaktadır, ancak uygulaması aynı değildir. Autograd tamamen Python’da yazılmıştır ve gradyanı doğrudan fonksiyondan hesaplarken, TensorFlow’un gradyan teyp işlevselliği ince bir Python sarıcı ile C++ ile yazılmıştır. TensorFlow, kayıptaki farklılıkları hesaplamak, kaybın gradyanını tahmin etmek ve bir sonraki en iyi adımı tahmin etmek için geri yayılımı kullanır.

XLA nedir?

XLA TensorFlow tarafından geliştirilen lineer cebir için alana özgü bir derleyicidir. TensorFlow belgelerine göre, XLA, potansiyel olarak kaynak kodu değişikliği olmadan TensorFlow modellerini hızlandırabilir, bu da hızı ve bellek kullanımını iyileştirir. Bir örnek 2020 Google’dır BERT MLPerf kıyaslama gönderimiburada XLA kullanan 8 Volta V100 GPU, ~7x performans iyileştirmesi ve ~5x parti boyutu iyileştirmesi elde etti.