前言

深度学习和人工智能引领了一个新的研究和发展方向，同时正在改变人类固有的处理问题的思维。现在各个领域都处于运用深度学习技术进行重大技术突破的阶段，与此同时，深度学习本身也展现出巨大的发展空间。

JAX是一个用于高性能数值计算的Python库，专门为深度学习领域的高性能计算而设计，其包含丰富的数值计算与科学计算函数，能够很好地满足用户的计算需求，特别是其基于GPU或者其他硬件加速器的能力，能够帮助我们在现有的条件下极大地加速深度学习模型的训练与预测。

JAX继承了Python简单易用的优点，给使用者提供了一个“便于入门，能够提高”的深度学习实现方案。JAX在代码结构上采用面向对象方法编写，完全模块化，并具有可扩展性，其运行机制和说明文档都将用户体验和使用难度纳入考虑范围，降低了复杂算法的实现难度。JAX的计算核心使用的是自动微分，可以支持自动模式反向传播和正向传播，且二者可以任意组合成任何顺序。

本书由浅到深地向读者介绍JAX框架相关的知识，重要内容均结合代码进行实战讲解，读者通过这些实例可以深入掌握JAX程序设计的内容，并能对深度学习有进一步的了解。

本书特色

版本新，易入门

本书详细介绍JAX最新版本的安装和使用，包括CPU版本以及GPU版本。

作者经验丰富，代码编写细腻

作者是长期奋战在科研和工业界的一线算法设计和程序编写人员，实战经验丰富，对代码中可能会出现的各种问题和“坑”有丰富的处理经验，使得读者能够少走很多弯路。

理论扎实，深入浅出

在代码设计的基础上，本书深入浅出地介绍深度学习需要掌握的一些基本理论知识，并通过大量的公式与图示对理论做介绍。

对比多种应用方案，实战案例丰富

本书给出了大量的实例，同时提供多个实现同类功能的解决方案，覆盖使用JAX进行深度学习开发中常用的知识。

本书内容

第1章　 JAX从零开始

本章介绍JAX应用于深度学习的基本理念、基础，并通过一个真实的深度学习例子向读者展示深度学习的一般训练步骤。本章是全书的基础，读者需要先根据本章内容搭建JAX开发环境，并下载合适的IDE。

第2章　 一学就会的线性回归、多层感知机与自动微分器

本章以深度学习中最常见的线性回归和多层感知机的程序设计为基础，循序渐进地介绍JAX进行深度学习程序设计的基本方法和步骤。

第3章　 深度学习的理论基础

本章主要介绍深度学习的理论基础，从BP神经网络开始，介绍神经网络两个基础算法，并着重介绍反向传播算法的完整过程和理论，最后通过编写基本Python的方式实现一个完整的反馈神经网络。

第4章　 XLA与JAX一般特性

本章主要介绍JAX的一些基础特性，例如XLA、自动微分等。读者需要了解的是XLA是如何工作的，它能给JAX带来什么。

第5章　 JAX的高级特性

本章是基于上一章的基础比较JAX与NumPy，重点解释JAX在实践中的一些程序设计和编写的规范要求，并对其中的循环函数做一个详尽而细致的说明。

第6章　 JAX的一些细节

本章主要介绍JAX在设计性能较优的程序时的细节问题，并介绍JAX内部一整套结构体保存方法和对模型参数的控制，这些都是为我们能编写出更为强大的深度学习代码打下基础。

第7章　 JAX中的卷积

卷积可以说是深度学习中使用最为广泛的计算部件，本章主要介绍卷积的基础知识以及相关用法，并通过一个经典的卷积神经网络VGG模型，讲解卷积的应用和JAX程序设计的一些基本内容。

第8章　 JAX与TensorFlow的比较与交互

本章主要介绍在一些需要的情况下使用已有的TensorFlow组件的一些方法。作为深度学习经典框架，TensorFlow有很多值得JAX参考和利用的内容。

第9章　 遵循JAX函数基本规则下的自定义函数

本章介绍JAX创建自定义函数的基本规则，并对其中涉及的一些细节问题进行详细讲解。期望读者在了解和掌握如何利用和遵循这些基本规则后去创建既满足需求又能够符合JAX函数规则的自定义函数。

第10章　 JAX中的高级包

本章详细介绍JAX中高级程序设计子包，特别是2个非常重要的模块jax.experimental和jax.nn。这两个包目前仍处于测试阶段，但是包含了建立深度学习模型所必需的一些基本函数。

第11章　 JAX实战—使用ResNet完成CIFAR100数据集分类

本章主要介绍在神经网络领域具有里程碑意义的模型—ResNet。它改变了人们仅仅依靠堆积神经网络层来获取更高性能的做法，在一定程度上解决了梯度消失和梯度爆炸的问题。这是一项跨时代的发明。本章以手把手的方式向读者介绍ResNet模型的编写和架构方法。

第12章　 JAX实战—有趣的词嵌入

本章介绍JAX于自然语言处理的应用，通过一个完整的例子向读者介绍自然语言处理所需要的全部内容，一步步地教会读者使用不同的架构和维度进行文本分类的方法。

第13章　 JAX实战—生成对抗网络（GAN）

本章介绍使用JAX完成生成对抗网络模型的设计，讲解如何利用JAX完成更为复杂的深度学习模型设计，掌握JAX程序设计的技巧。同时也期望通过本章能够帮助读者全面复习本书所涉及的JAX的深度学习程序设计内容。

源码下载与技术支持

本书配套源码、PPT课件、数据集、开发环境、配图文件和答疑服务，需要使用微信扫描右边二维码下载，可按页面提示，把链接转发到自己的邮箱中下载。如果下载有问题或者阅读中发现问题，请联系booksaga@163.com，邮件主题为“谷歌JAX深度学习从零开始学”。

本书读者

●　人工智能入门读者

●　深度学习入门读者

●　机器学习入门读者

●　高等院校人工智能专业的师生

●　专业培训机构的师生

●　其他对智能化、自动化感兴趣的开发者

技术支持、勘误和鸣谢

由于作者的水平有限，加上JAX框架的演进较快，书中难免存在疏漏之处，恳请读者来信批评指正。本书的顺利出版，首先要感谢家人的理解和支持，他们给予我莫大的动力，让我的努力更加有意义。此外特别感谢出版社的编辑们，感谢他们在本书编写过程中给予的无私指导。

编者

2022年4月