Название: Deep Learning with JAX (Final Release)
Автор: Grigory Sapunov
Издательство: Manning Publications
Год: 2024
Страниц: 410
Язык: английский
Формат: pdf (true)
Размер: 39.8 MB

Accelerate Deep Learning and other number-intensive tasks with JAX, Google’s awesome high-performance numerical computing library.

In Deep Learning with JAX you will learn how to:

Use JAX for numerical calculations
Build differentiable models with JAX primitives
Run distributed and parallelized computations with JAX
Use high-level neural network libraries such as Flax and Haiku
Leverage libraries and modules from the JAX ecosystem

The JAX numerical computing library tackles the core performance challenges at the heart of Deep Learning and other scientific computing tasks. By combining Google’s Accelerated Linear Algebra platform (XLA) with a hyper-optimized version of NumPy and a variety of other high-performance features, JAX delivers a huge performance boost in low-level computations and transformations.

Deep Learning with JAX is a hands-on guide to using JAX for Deep Learning and other mathematically-intensive applications. Google Developer Expert Grigory Sapunov steadily builds your understanding of JAX’s concepts. The engaging examples introduce the fundamental concepts on which JAX relies and then show you how to apply them to real-world tasks. You’ll learn how to use JAX’s ecosystem of high-level libraries and modules, and also how to combine TensorFlow and PyTorch with JAX for data loading and deployment.

About the technology

The JAX Python mathematics library is used by many successful Deep Learning organizations, including Google’s groundbreaking DeepMind team. This exciting newcomer already boasts an amazing ecosystem of tools including high-level Deep Learning libraries Flax by Google, Haiku by DeepMind, gradient processing and optimization libraries, libraries for evolutionary computations, Federated Learning, and much more! JAX brings a functional programming mindset to Python Deep Learning, letting you improve your composability and parallelization in a cluster.

About the book:

Deep Learning with JAX teaches you how to use JAX and its ecosystem to build neural networks. You’ll learn by exploring interesting examples including an image classification tool, an image filter application, and a massive scale neural network with distributed training across a cluster of TPUs. Discover how to work with JAX for hardware and other low-level aspects and how to solve common Machine Learning problems with JAX. By the time you’re finished with this awesome book, you’ll be ready to start applying JAX to your own research and prototyping!

Chapter 1 introduces you to the core concepts and strengths of JAX. It sets the stage by delving into the history of deep learning frameworks, highlighting the evolutionary path that led to the birth of JAX. This chapter illuminates JAX’s niche in the ecosystem, showcasing its adeptness at compiling and parallelizing code for a wide range of applications, from ocean simulation to large-scale neural networks. By comparing JAX with TensorFlow, PyTorch, and NumPy, you’ll understand the strategic scenarios that call for JAX’s utilization, laying a robust foundation for the rest of the book.

Chapter 2 is a hands-on exploration that walks you through the development of a simple neural network application. This chapter is your gateway to experiencing JAX’s transformative capabilities firsthand. You’ll learn about the high-level structure of a JAX project, from loading datasets to creating neural networks, and discover the power of JAX transformations. The chapter concludes with a dive into saving and loading models and the distinction between pure and impure functions, equipping you with the knowledge to embark on more complex projects.

Chapters 3 through 10 are meticulously designed to guide you through the intricacies of JAX, ensuring a mastery of its most powerful features. You’ll start by working with arrays (chapter 3), to understand the nuances that differentiate JAX from NumPy and learn how to use these differences to your advantage. As you progress, you’ll delve into calculating gradients (chapter 4) using JAX’s automatic differentiation capabilities to simplify and accelerate the training of neural networks. Chapters 5 and 6 introduce you to JAX’s just-in-time compilation and auto-vectorization, revealing strategies to significantly boost performance.

The journey through Core JAX doesn’t stop there. You’ll explore parallelizing your computations and using tensor sharding (chapters 7 and 8) to scale your applications across multiple devices, a critical skill for handling large-scale models and datasets. The exploration of random numbers in JAX (chapter 9) unveils the functional approach to randomness, ensuring reproducibility and efficiency in stochastic operations. Finally, chapter 10, on working with pytrees, teaches you to manage complex data structures elegantly, further enhancing your ability to build sophisticated models and algorithms with JAX.

Chapter 11 discusses higher-level neural network libraries, focusing on Flax and its Linen API, along with Optax for gradient transformations. You’ll learn how to construct models more intuitively, manage training states, and interact with the Hugging Face ecosystem to access state-of-the-art models. This chapter bridges the gap between JAX’s core capabilities and the practical needs of Deep Learning projects, providing you with the tools to build, train, and deploy complex models efficiently.

Chapter 12 takes a broader look at the JAX ecosystem, showcasing libraries for various Machine Learning tasks, including reinforcement learning and evolutionary computations. We’ll also explore JAX modules for other scientific fields like physics, chemistry, and more.

About the reader:
For intermediate Python programmers who are familiar with Deep Learning.

About the author:
Grigory Sapunov is a co-founder and CTO of Intento. He is a software engineer with more than twenty years of experience. Grigory holds a Ph.D. in artificial intelligence and is a Google Developer Expert in Machine Learning.