用JAX构建高性能神经网络 · 课程目录

为什么选择JAX？NumPy关系 · 函数式设计 · 安装配置

JAX基础操作

数组创建 · NumPy互操作 · 随机数种子

JAX核心机制

JIT编译 · vmap自动向量化 · pmap · grad

JAX函数式编程

纯函数 · 状态管理 · 参数更新模式

JAX神经网络基础

stax构建网络 · 参数初始化 · 前向传播

JAX优化器实现

梯度下降 · SGD动量 · Adam从零实现

JAX训练循环

手动训练循环 · 损失函数 · 批次加载

JAX与Flax框架

Flax简介 · Module系统 · 协同工作

Flax高级模块

自定义Module · 参数共享 · 检查点

JAX数据加载

tf.data配合 · 自定义管道 · 预处理增强

JAX性能优化

内存管理 · shard_map · 编译缓存

JAX分布式训练

数据并行 · 模型并行 · 流水线并行

JAX混合精度训练

float16/bfloat16 · 自动混合精度 · 损失缩放

JAX自定义梯度

自定义VJP · 梯度裁剪 · 梯度累积

JAX高级自动微分

Hessian · Jacobian · 高阶梯度

JAX概率编程

NumPyro · 概率模型 · 变分推断

JAX强化学习应用

环境交互 · 策略梯度 · 价值函数

JAX自然语言处理

Transformer实现 · 注意力 · 序列生成

JAX计算机视觉

卷积网络 · 图像分类 · 迁移学习

JAX生成模型

VAE · GAN训练 · 扩散模型入门

JAX图神经网络

消息传递 · 图卷积 · 分子预测

JAX时间序列分析

RNN/LSTM · 序列预测 · 注意力时间模型

JAX模型部署

SavedModel · TFLite · ONNX导出

JAX与TensorBoard

日志记录 · 可视化曲线 · 超参数调优

JAX调试技巧

pdb调试JIT · 打印中间值 · 性能分析

JAX生态工具

Optax · Haiku · RLax

JAX大规模训练

TPU指南 · 多主机 · 梯度压缩

JAX最佳实践

代码组织 · 测试 · 版本控制与复现

JAX前沿研究

隐式微分 · 神经ODE · 可微编程

JAX综合项目

端到端流水线 · 评估部署 · 总结展望