JAX是一个用于高性能数值计算的Python库，专门为深度学习领域的高性能计算而设计。本书详解JAX框架深度学习的相关知识，配套示例源码、PPT课件、数据集和开发环境。 本书共分为13章，内容包括JAX从零开始，一学就会的线性回归、多层感知机与自动微分器，深度学习的理论基础，XLA与JAX一般特性，JAX的高级特性，JAX的一些细节，JAX中的卷积，JAX与TensorFlow的比较与交互，遵循JAX函数基本规则下的自定义函数，JAX中的高级包。后给出3个实战案例：使用ResNet完成CIFAR100数据集分类，有趣的词嵌入，生成对抗网络（GAN）。 本书适合JAX框架初学者、深度学习初学者以及深度学习从业人员，也适合作为高等院校和培训机构人工智能相关专业的师生教学参考书。

JAX是一个用于高性能数值计算的Python库，专门为深度学习领域的高性能计算而设计。本书详解JAX框架深度学习的相关知识，并剖析3个实战案例： 使用ResNet完成CIFAR100数据集分类、有趣的词嵌入与生成对抗网络。本书配套示例源码、PPT课件、数据集、开发环境和答疑服务。

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

王晓华，计算机专业讲师，研究方向为云计算、大数据与人工智能。著有《Spark MLlib机器学习实践》《TensorFlow深度学习应用实践》《OpenCV TensorFlow深度学习与计算机视觉实战》《TensorFlow知识图谱实战》《TensorFlow人脸识别实战》《TensorFlow语音识别实战》《TensorFlow 2.0卷积神经网络实战》《Keras实战：基于TensorFlow2.2的深度学习实践》《TensorFlow深度学习从零开始学》《深度学习的数学原理与实现》等图书。

第 1 章 JAX从零开始 1
1.1 JAX来了 1
1.1.1 JAX是什么 1
1.1.2 为什么是JAX 2
1.2 JAX的安装与使用 3
1.2.1 Windows Subsystem for Linux的安装 3
1.2.2 JAX的安装和验证 7
1.2.3 PyCharm的下载与安装 8
1.2.4 使用PyCharm和JAX 9
1.2.5 JAX的Python代码小练习：计算SeLU函数 11
1.3 JAX实战MNIST手写体的识别 12
1.3.1 第一步：准备数据集 12
1.3.2 第二步：模型的设计 13
1.3.3 第三步：模型的训练 13
1.4 本章小结 15
第2章 一学就会的线性回归、多层感知机与自动微分器 16
2.1 多层感知机 16
2.1.1 全连接层多层感知机的隐藏层 16
2.1.2 使用JAX实现一个全连接层 17
2.1.3 更多功能的全连接函数 19
2.2 JAX实战鸢尾花分类 22
2.2.1 鸢尾花数据准备与分析 23
2.2.2 模型分析采用线性回归实战鸢尾花分类 24
2.2.3 基于JAX的线性回归模型的编写 25
2.2.4 多层感知机与神经网络 27
2.2.5 基于JAX的激活函数、softmax函数与交叉熵函数 29
2.2.6 基于多层感知机的鸢尾花分类实战 31
2.3 自动微分器 35
2.3.1 什么是微分器 36
2.3.2 JAX中的自动微分 37
2.4 本章小结 38
第3章 深度学习的理论基础 39
3.1 BP神经网络简介 39
3.2 BP神经网络两个基础算法详解 42
3.2.1 最小二乘法详解 43
3.2.2 道士下山的故事梯度下降算法 44
3.2.3 最小二乘法的梯度下降算法以及JAX实现 46
3.3 反馈神经网络反向传播算法介绍 52
3.3.1 深度学习基础 52
3.3.2 链式求导法则 53
3.3.3 反馈神经网络原理与公式推导 54
3.3.4 反馈神经网络原理的激活函数 59
3.3.5 反馈神经网络原理的Python实现 60
3.4 本章小结 64
第4章 XLA与JAX一般特性 65
4.1 JAX与XLA 65
4.1.1 XLA如何运行 65
4.1.2 XLA如何工作 67
4.2 JAX一般特性 67
4.2.1 利用JIT加快程序运行 67
4.2.2 自动微分器grad函数 68
4.2.3 自动向量化映射vmap函数 70
4.3 本章小结 71
第5章 JAX的高级特性 72
5.1 JAX与NumPy 72
5.1.1 像NumPy一样运行的JAX 72
5.1.2 JAX的底层实现lax 74
5.1.3 并行化的JIT机制与不适合使用JIT的情景 75
5.1.4 JIT的参数详解 77
5.2 JAX程序的编写规范要求 78
5.2.1 JAX函数必须要为纯函数 79
5.2.2 JAX中数组的规范操作 80
5.2.3 JIT中的控制分支 83
5.2.4 JAX中的if、while、for、scan函数 85
5.3 本章小结 89
第6章 JAX的一些细节 90
6.1 JAX中的数值计算 90
6.1.1 JAX中的grad函数使用细节 90
6.1.2 不要编写带有副作用的代码JAX与NumPy的差异 93
6.1.3 一个简单的线性回归方程拟合 94
6.2 JAX中的性能提高 98
6.2.1 JIT的转换过程 98
6.2.2 JIT无法对非确定参数追踪 100
6.2.3 理解JAX中的预编译与缓存 102
6.3 JAX中的函数自动打包器vmap 102
6.3.1 剥洋葱对数据的手工打包 102
6.3.2 剥甘蓝JAX中的自动向量化函数vmap 104
6.3.3 JAX中高阶导数的处理 105
6.4 JAX中的结构体保存方法Pytrees 106
6.4.1 Pytrees是什么 106
6.4.2 常见的pytree函数 107
6.4.3 深度学习模型参数的控制（线性模型） 108
6.4.4 深度学习模型参数的控制（非线性模型） 113
6.4.5 自定义的Pytree节点 113
6.4.6 JAX数值计算的运行机制 115
6.5 本章小结 117
第7章 JAX中的卷积 118
7.1 什么是卷积 118
7.1.1 卷积运算 119
7.1.2 JAX中的一维卷积与多维卷积的计算 120
7.1.3 JAX.lax中的一般卷积的计算与表示 122
7.2 JAX实战基于VGG架构的MNIST数据集分类 124
7.2.1 深度学习Visual Geometry Group（VGG）架构 124
7.2.2 VGG中使用的组件介绍与实现 126
7.2.3 基于VGG6的MNIST数据集分类实战 129
7.3 本章小结 133
第8章 JAX与TensorFlow的比较与交互 134
8.1 基于TensorFlow的MNIST分类 134
8.2 TensorFlow与JAX的交互 137
8.2.1 基于JAX的TensorFlow Datasets数据集分类实战 137
8.2.2 TensorFlow Datasets数据集库简介 141
8.3 本章小结 145
第9章 遵循JAX函数基本规则下的自定义函数 146
9.1 JAX函数的基本规则 146
9.1.1 使用已有的原语 146
9.1.2 自定义的JVP以及反向VJP 147
9.1.3 进阶jax.custom_jvp和jax.custom_vjp函数用法 150
9.2 Jaxpr解释器的使用 153
9.2.1 Jaxpr tracer 153
9.2.2 自定义的可以被Jaxpr跟踪的函数 155
9.3 JAX维度名称的使用 157
9.3.1 JAX的维度名称 157
9.3.2 自定义JAX中的向量Tensor 158
9.4 本章小结 159
第10章 JAX中的高级包 160
10.1 JAX中的包 160
10.1.1 jax.numpy的使用 161
10.1.2 jax.nn的使用 162
10.2 jax.experimental包和jax.example_libraries的使用 163
10.2.1 jax.experimental.sparse的使用 163
10.2.2 jax.experimental.optimizers模块的使用 166
10.2.3 jax.experimental.stax的使用 168
10.3 本章小结 168
第11章 JAX实战使用ResNet完成CIFAR100数据集分类 169
11.1 ResNet基础原理与程序设计基础 169
11.1.1 ResNet诞生的背景 170
11.1.2 使用JAX中实现的部件不要重复造轮子 173
11.1.3 一些stax模块中特有的类 175
11.2 ResNet实战CIFAR100数据集分类 176
11.2.1 CIFAR100数据集简介 176
11.2.2 ResNet残差模块的实现 179
11.2.3 ResNet网络的实现 181
11.2.4 使用ResNet对CIFAR100数据集进行分类 182
11.3 本章小结 184
第12章 JAX实战有趣的词嵌入 185
12.1 文本数据处理 185
12.1.1 数据集和数据清洗 185
12.1.2 停用词的使用 188
12.1.3 词向量训练模型word2vec的使用 190
12.1.4 文本主题的提取：基于TF-IDF 193
12.1.5 文本主题的提取：基于TextRank 197
12.2 更多的词嵌入方法FastText和预训练词向量 200
12.2.1 FastText的原理与基础算法 201
12.2.2 FastText训练以及与JAX的协同使用 202
12.2.3 使用其他预训练参数嵌入矩阵（中文） 204
12.3 针对文本的卷积神经网络模型字符卷积 205
12.3.1 字符（非单词）文本的处理 206
12.3.2 卷积神经网络文本分类模型的实现conv1d（一维卷积） 213
12.4 针对文本的卷积神经网络模型词卷积 216
12.4.1 单词的文本处理 216
12.4.2 卷积神经网络文本分类模型的实现 218
12.5 使用卷积对文本分类的补充内容 219
12.5.1 中文的文本处理 219
12.5.2 其他细节 222
12.6 本章小结 222
第13章 JAX实战生成对抗网络（GAN） 223
13.1 GAN的工作原理详解 223
13.1.1 生成器与判别器共同构成了一个GAN 224
13.1.2 GAN是怎么工作的 225
13.2 GAN的数学原理详解 225
13.2.1 GAN的损失函数 226
13.2.2 生成器的产生分布的数学原理相对熵简介 226
13.3 JAX实战GAN网络 227
13.3.1 生成对抗网络GAN的实现 228
13.3.2 GAN的应用前景 232
13.4 本章小结 235
附录 Windows 11安装GPU版本的JAX 236