کتابخانه JAX برای یادگیری ماشین: چگونه کار می کند و چرا آن را یاد بگیریم

کتابخانه Jax که توسط تیم Google Research معرفی شده هر روز محبوبیت بیشتری پیدا می کند.کمپانی Deepmind اخیراً اعلام کرد که از JAX برای تسریع تحقیقات خود استفاده می کند و در حال حاضر مجموعه ای از کتابخانه ها را بر روی JAX توسعه داده است.مقالات تحقیقاتی بیشتری با استفاده از JAX ساخته شده است، به عنوان مثال، Vision Transformer (ViT).کتابخانه JAX عملکرد چشمگیری را نشان می دهد. و حتی سریعترین ترانسفورماتور جهان هم اکنون بر روی JAX ساخته شده است.موفقیت اخیر JAX مدل GPT-J-6B توسط EleutherAI است. به گفته نویسندگان: "این پروژه به میزان قابل توجهی فرد-ساعت کمتری در مقایسه با سایر مدل های مقیاس بزرگ نیاز داشت، که نشان می دهد JAX + xmap + TPU مجموعه مناسبی از ابزارها برای توسعه سریع مدل های بزرگ است."جکس (JAX) چیست؟و نقاط قوت آن چیست؟جکس (JAX) یک کتابخانه بسیار سطح پایین شبیه NumPy است اما دارای چندین ویژگی جالب است:اتوگراد (Autograd)- JAX می تواند به طور خودکار کد بومی Python و NumPy را دیفرانسیل بگیرد. همچنین می تواند از طریق زیرمجموعه وسیعی از ویژگی های پایتون، از جمله حلقه ها، ifها، بازگشت، دیفرانسیل بگیرد و حتی می تواند مشتقِ مشتقِ مشتق را نیز بگیرد. همچنین از reverse-mode و forward-mode دیفرانسیل پشتیبانی می کند، و این دو را می توان به دلخواه به هر ترتیبی ترکیب کند. کتابخانه قدیمی خوب autograd پدر JAX است، و برخی از نویسندگان autograd در حال توسعه JAX هستند. اگر به دنبال درک بهتری از دیفرانسیل خودکار، تفاوت دیفرانسیل عددی یا نمادین، تفاوت reverse-mode و forward-mode و غیره هستید، پس این مقاله خوب را در مورد این موضوع بخوانید.کامپایل کد شما- JAX از XLA برای کامپایل و اجرای کد NumPy شما روی شتاب دهنده ها مانند GPU ها و TPU ها استفاده می کند. کامپایل به طور پیش فرض در زیر پوسته اتفاق می افتد و تماس های کتابخانه به just-in-time (JIT) کامپایل و اجرا می شوند. جیت(JIT)- JAX حتی به شما امکان می دهد تا به موقع (just-in-time) توابع Python خود را با استفاده از یک API تک کاره در هسته های بهینه شده با XLA کامپایل کنید. گردآوری (Compilation) و دیفرانسیل خودکار را می توان به صورت دلخواه ترکیب کرد، بنابراین می توانید الگوریتم های پیچیده را بیان کرده و حداکثر کارایی را بدون نیاز به خروج از پایتون داشته باشید.گسیل غیر همزمان (Asynchronous dispatch) مفید است زیرا به پایتون اجازه می دهد تا از دستگاه شتاب دهنده "run ahead" اجرا شود و کد پایتون را از مسیر بحرانی دور نگه دارد. به شرطی که کد پایتون سریعتر از اجرای دستگاه کار کند و به شرطی که کد پایتون در واقع نیازی به بررسی خروجی محاسبه بر روی هاست نداشته باشد، یک برنامه پایتون می تواند مقادیر دلخواه کار را Enqueued کند و از منتظر ماندن شتاب دهنده اجتناب کنید. بردارسازی اوتوماتيک به ويمپ (Auto-vectorization with vmap) این دارای معانی آشنایی برای نگاشتن یک تابع در امتداد محورهای آرایه است، اما به جای نگه داشتن حلقه در خارج ، حلقه را برای عملکرد بهتر یک تابع به سمت پایین هدایت می کند. هنگامی که با JIT ترکیب می شود، می تواند به همان اندازه سریع باشد که ابعاد دسته را با دست اضافه کنید.موازی سازی با پی مپ (Parallelization with pmap.) هدف ()pmap بیان برنامه های single-program multiple-data (SPMD) است. اعمال ()pmap بر روی یک تابع، عملکرد را با XLA (مشابه() jit ) کامپایل می کند، سپس آن را به طور موازی در دستگاه های XLA، مانند GPU های متعدد یا چندین هسته TPU اجرا می کند. از نظر معنایی قابل مقایسه با () vmap است زیرا هر دو تبدیل یک تابع را روی محورهای آرایه ترسیم می کنند، اما در جایی که ()vmap با pushing محور نگاشت شده به سمت عملیات ابتدایی توابع را برداری می کند،() pmap در عوض تابع را تکرار می کند و هر کپی را بر روی دستگاه XLA خود به طور موازی اجرا می کند.مدل برنامه نویسی محور نامگذاری شده با xmap (آزمایشی)(Named-axis programming model with xmap) این به شما کمک می کند تا توابع اجتناب از خطا و خود مستندسازی(self-documenting) را با استفاده از محورهای نامگذاری بنویسید، سپس نحوه اجرای آنها را در سخت افزار در هر مقیاسی، از CPU لپ تاپ تا بزرگترین ابر رایانه TPU، کنترل کنید. killer feature xmap توانایی آن برای موازی سازی کد بر روی سخت افزارهای مقیاس سوپ کامپیوترها است!در اینجا یک معرفی کوتاه برای JAX (7 دقیقه) است:https://youtu.be/0mVmRHMaOJ4و یک مقدمه عمیق تر (23 دقیقه):https://youtu.be/z-WSrQDXkuMو حتی عمیق تر (45 دقیقه):https://developer.nvidia.com/gtc/2020/video/s21989و عمیق ترین اون ها که از یکی نویسندگان JAX (1h 10m) است:https://youtu.be/BzuEGdGHKjcو چند مقاله در JAX:یک معرفی کاربردی خوب توسط Simone Scardapane: “JAX, AKA NUMPY ON STEROIDS”.“From PyTorch to JAX: towards neural net frameworks that purify stateful code”نوعی مقایسه عملی رویکردهای JAX و PyTorch است.JAX vs PyTorchنزدیکترین چارچوب یادگیری ماشین به JAX ،PyTorch است. JAX با تعاریف عملکرد سطح پایین تر برای کارهای تحقیقاتی مناسب است. با این اوصاف، PyTorch وسعت بسیار بیشتری از کتابخانه ها و ابزارهای کاربردی، تعاریف شبکه از پیش آموزش دیده و از پیش نوشته شده، بارگذار داده و قابلیت حمل در مقاصد استقرار را ارائه می دهد.JAX vs TensorFlowهر دوی JAX و TensorFlow توسط Google نوشته شده اند. از تجربیات اولیه، توسعه JAX بسیار ساده تر به نظر می رسد و بصری تر است. با این اوصاف، JAX فاقد زیرساخت های گسترده ای است که TensorFlow در طول این سالها ایجاد کرده است خواه پروژه های منبع باز، مدل های از پیش آموزش دیده، آموزش ها، چکیده های سطح بالاتر (از طریق Keras) و قابلیت حمل در مقاصد استقرار.کی به سراغ JAX برویمکتابخانه JAX یک چارچوب یادگیری ماشین جدید است که در تحقیقات یادگیری ماشین محبوبیت زیادی پیدا کرده است.اگر در حوزه تحقیقاتی فعالیت می کنید، JAX گزینه خوبی برای پروژه شما است.اگر به طور فعال برنامه ای را توسعه می دهید، چارچوب های PyTorch و TensorFlow ابتکار شما را با سرعت بیشتر پيش خواهد برد. و البته، در بینایی رایانه ای همیشه در مقایسه با خرید ابزار بینایی رایانه ای و ساخت آن یک مصالحه ای وجود دارد.کتابخانه های بالای JAXچندین کتابخانه سطح بالا برای شبکه های عصبی در بالای JAX وجود دارد:کتابخانه Flax (Google Brain) — یک کتابخانه شبکه عصبی سطح بالا که برای انعطاف پذیری طراحی شده است (در مقایسه با Keras ، Sonnet یا Haiku). Repo: https://github.com/google/flaxکتاب خانه Haiku (Deepmind) — یک کتابخانه شبکه عصبی مبتنی بر JAX، معروف به Sonnet for JAX (اگر نمی دانید Sonnet چیست، شبیه کتابخانه شبکه عصبی سطح بالای Keras در بالای TensorFlow است).Repo: https://github.com/deepmind/dm-haikuکتاب خانه Trax (Google Brain) — یک کتابخانه end-to-end برای یادگیری عمیق که بر کد واضح و سرعت تمرکز دارد. Trax شامل مدل های اصلی (مانند الگوریتم های ResNet, LSTM, Transformer و RL (مانند REINFORCE, A2C, PPO) است. همچنین از آن به طور فعال برای تحقیقات استفاده می شود و شامل مدل های جدیدی مانند Reformer و الگوریتم های RL جدید مانند AWR است. Trax دارای پیوندهایی به تعداد زیادی مجموعه داده یادگیری عمیق، از جمله مجموعه داده های Tensor2Tensor و TensorFlow datasets است.Repo: https://github.com/google/traxکتابخانه Objax (Google) — یک چارچوب شی گرا مینیمالیستی با رابط کاربری شبیه PyTorch. نام آن از ادغام Object و JAX گرفته شده است. Objax توسط محققان با تمرکز بر سادگی و قابل فهم بودن توسط محققان طراحی شده است. کاربران آن باید بتوانند آن را به راحتی بخوانند، بفهمند، بسط دهند و متناسب با نیازهای خود تغییر دهند.Repo: https://github.com/google/objaxکتابخانه Stax (experimental part of JAX) — یک کتابخانه مشخصات شبکه عصبی کوچک اما انعطاف پذیر از پایه.Repo: https://jax.readthedocs.io/en/latest/jax.experimental.stax.htmlکتابخانه Elegy (PoetsAI) — یک چارچوب شبکه های عصبی بر اساس Jax و با الهام ازKeras. چهارچوب Elegy رابط Keras API را پیاده سازی می کند اما نقش بهتری را در Jax بازی می کند، انعطاف پذیری بیشتری در مورد losses and metrics و module system ایجاد می کند که استفاده از آن را بسیار آسان می کند.Repo: https://github.com/poets-ai/elegyکتابخانه Mesh Transformer JAX — چارچوبی که برای ساخت مدل زبان GPT-J-6B اخیر استفاده شد. این مدل ترانسفورماتورهای موازی را در JAX و Haiku ارائه می دهد. این نرم افزار برای مقیاس پذیری تا حدود 20B پارامتر در TPUv3 طراحی شده است، که فراتر از آن باید از استراتژی های مختلف موازی سازی استفاده شود.Repo: https://github.com/kingoflolz/mesh-transformer-jax‌کتابخانه Swarm JAX — چارچوب آموزشی Swarm با استفاده از Haiku + JAX + Ray برای مدل های زبان ترانسفورماتور موازی لایه در گره های نامعتبر و ناهمگن است.Repo: https://github.com/kingoflolz/swarm-jaxچند کتابخانه RL (اما برخی از کتابخانه های سطح بالا، مثلاً Trax، دارای ویژگی های RL نیز هستند):کتابخانه RLax (Deepmind) — بلوک های سازنده برای اجرای عوامل یادگیری تقویتی اجزای موجود در RLax طیف وسیعی از الگوریتم ها و ایده ها را شامل می شود: یادگیری TD، policy gradients،actor critics، MAP ،proximal policy optimization ،non-linear value transformation ،non-linear value transformation general value functions و تعدادی روش کاوش.Repo: https://github.com/deepmind/rlaxکتابخانه Coax (Microsoft) — یک بسته پایتون یادگیری تقویتی مدولار (RL) برای حل محیط های OpenAI Gym با تقریب عملکردها بر اساس JAX است.Repo: https://github.com/microsoft/coaxبرخی دیگر از کتابخانه های تخصصی تر:کتابخانه Optax (Deepmind) — یک کتابخانه پردازش گرادیان و بهینه سازی برای JAX. چهارچوب Optax یک کتابخانه از تغییرات گرادیان، همراه با عملگرهای ترکیب (به عنوان مثال زنجیره) را فراهم می کند که امکان اجرای بسیاری از بهینه سازهای استاندارد (مانند RMSProp یا Adam) را در یک خط کد فراهم می کند.Repo: https://github.com/deepmind/optaxکتابخانه Chex (Deepmind) — کتابخانه ای از ابزارها برای کمک به نوشتن کد JAX قابل اعتماد. Chex مجموعه ای از ابزارها را شامل JAX-aware unit testing،assertions of properties of JAX datatypes, mocks and fakes و محیط های آزمایش چند دستگاه را ارائه می دهد.Repo: https://github.com/deepmind/chexکتابخانه Jraph, که “giraffe” تلفظ می شود, (Deepmind) — یک کتابخانه سبک برای کار با شبکه های عصبی گرافیکی در JAX. کتابخانه Jraph یک ساختار داده استاندارد برای نمودارها، مجموعه ای از ابزارهای مفید برای کار با نمودارها و یک "باغ وحش" از مدل های شبکه عصبی گرافیکی به راحتی قابل حمل و قابل توسعه را فراهم می کند.Repo: https://github.com/deepmind/jraphکتابخانه JAX, M.D. (Google) — چارچوبی برای Differentiable Physics.که امکانتی همچون differentiable، شتاب دهنده سخت افزاری، دینامیک مولکولی را بر اساس JAX فراهم می کند.Repo: https://github.com/google/jax-mdکتابخانه Oryx (Google?) کتابخانه ای برای برنامه نویسی احتمالی و یادگیری عمیق است که در بالای جکس ساخته شده است. Oryx یک کتابخانه تجربی است که JAX را به برنامه های کاربردی اعم از ساخت و آموزش شبکه های عصبی پیچیده تا تقریب استنباط بیزی در مدل های generative عمیق گسترش می دهد. مانند JAX که jit ،vmap و grad را ارائه می دهد، Oryx مجموعه ای از تغییرات عملکرد ترکیبی را ارائه می دهد که نوشتن کد ساده و تبدیل آن به پیچیدگی را ایجاد می کند در حالی که کاملاً با JAX همکاری می کند.Repo: https://github.com/tensorflow/probability/tree/master/spinoffs/oryxمنابع:https://moocaholic.medium.com/jax-a13e83f49897https://blog.roboflow.com/jax-framework/