📘 JAX 梯度累积 & 混合精度
30 章 · 实战课程
🎯 友好色系 · 紧凑目录
01
JAX基础回顾
JAX是什么
与NumPy
jit/vmap/pmap
安装配置
02
梯度累积原理
为什么需要
数学原理
Batch Size关系
适用场景
03
JAX梯度计算
grad详解
value_and_grad
高阶梯度
注意事项
04
手动实现梯度累积
for循环
lax.scan
代码模板
调试技巧
05
JAX优化器与梯度累积
optax介绍
optax结合
MultiSteps
参数配置
06
混合精度训练基础
FP16/FP32
原理
NVIDIA AMP
硬件要求
07
JAX混合精度
jax.numpy精度
convert_element_type
默认策略
自定义
08
JAX AMP实战
auto_sharding
jax.experimental.amp
loss scaling
梯度裁剪
09
梯度累积+混合精度组合
必要性
注意事项
性能分析
常见错误
10
实战：图像分类 (ResNet-18)
CIFAR-10
ResNet-18
训练循环
梯度累积+混合精度
11
实战：NLP模型 (Transformer)
Transformer简介
文本预处理
训练循环
梯度累积+混合精度
12
性能调优
内存分析
速度优化
步数选择
精度选择
13
分布式训练与梯度累积
pmap
分布式原理
实现
同步策略
14
梯度累积的变体
梯度检查点
梯度压缩
异步更新
局部更新
15
混合精度进阶技巧
动态/静态loss scaling
+BN
+LN
16
JAX自定义梯度
custom_jvp/vjp
+混合精度
+梯度累积
调试
17
梯度累积与模型并行
模型并行概念
应用
数据并行结合
梯度同步
18
混合精度与模型量化
量化概念
混合精度+量化
JAX量化工具
量化感知训练
19
实战：GAN训练
GAN原理
生成器/判别器
梯度累积应用
混合精度应用
20
实战：强化学习
基本概念
策略梯度
梯度累积应用
混合精度应用
21
梯度累积调试与可视化
梯度范数监控
直方图
步数选择策略
收敛性分析
22
混合精度训练调试
精度溢出检测
NaN/Inf处理
loss scaling调试
收敛性分析
23
JAX自动微分进阶
grad底层
vjp/jvp
linearize
custom_vjp高级
24
梯度累积与超参数搜索
学习率关系
Batch Size关系
搜索策略
贝叶斯优化
25
混合精度与超参数搜索
精度策略&LR
精度&Batch Size
精度选择
自动化
26
JAX生态：梯度累积
Flax
Haiku
Trax
其他库
27
JAX生态：混合精度
Flax混合精度
Haiku
Trax
其他库
28
TPU实践
TPU架构
梯度累积
混合精度
性能优化
29
GPU实践
GPU架构
梯度累积
混合精度
性能优化
30
总结与展望
最佳实践
常见问题
未来趋势
学习资源