JAX 与 PyTorch 模型互转实战 · 课程目录

🏆 动态图 ↔ 静态图 · 函数式 ↔ 面向对象 📁 01.html ~ 30.html 一键跳转

JAX与PyTorch生态全景对比核心差异

为什么需要互转？动态图vs静态图、函数式vs面向对象、适用场景分析。

环境准备与工具链搭建安装

JAX、PyTorch、Flax、Haiku、Trax等库安装与版本兼容性检查。

PyTorch模型基础回顾 nn.Module

forward方法、state_dict、模型保存与加载 (torch.save/load)。

JAX模型基础入门 Flax/Haiku

Flax的nn.Module、Haiku的hk.Module、JAX纯函数式编程范式。

参数结构映射 state_dict ⇄ params

PyTorch state_dict 与 JAX/Flax 参数字典结构对比与转换。

线性层 (Linear) 互转 Dense

PyTorch nn.Linear → Flax nn.Dense，权重与偏置维度对齐。

卷积层 (Conv2d) 互转 padding/stride

PyTorch nn.Conv2d ↔ Flax nn.Conv，padding与stride差异处理。

批归一化 (BatchNorm) 互转 running stats

track_running_stats ↔ use_running_average，推理模式切换。

激活函数与Dropout互转无状态层

ReLU、GELU、Dropout等直接映射。

Embedding层互转权重共享

PyTorch nn.Embedding ↔ Flax nn.Embed，权重共享处理。

序列模型 (RNN/LSTM/GRU) 互转 LSTMCell

PyTorch nn.LSTM ↔ Flax nn.LSTMCell，单向与双向适配。

Transformer核心组件互转 Attention

Self-Attention、Multi-Head Attention 的 PyTorch/Flax 实现对比。

优化器状态互转 Optax

optimizer.state_dict ↔ Optax 状态，学习率调度器迁移。

手动权重搬运实战 numpy中间格式

逐层复制参数，验证数值一致性。

自动转换工具开发递归遍历

Python脚本自动生成转换映射表。

Flax模型转PyTorch params → state_dict

从Flax参数字典重建PyTorch state_dict，处理命名差异。

PyTorch模型转Flax state_dict → params

生成Flax参数字典，处理BatchNorm running stats。

Haiku模型与PyTorch互转 hk.Params

Haiku的hk.Params ↔ state_dict 转换技巧。

JAX原生函数式模型转PyTorch 范式冲突

处理JAX纯函数与PyTorch nn.Module 的范式冲突。

模型推理一致性验证数值误差分析

相同随机输入，比较PyTorch与JAX输出差异。

梯度与反向传播对比 autograd vs grad

PyTorch autograd ↔ JAX grad/vjp，验证梯度数值一致性。

Batch维度处理差异 batch_first

PyTorch默认batch_first，JAX/Flax batch维度位置，reshape注意。

设备管理 (CPU/GPU/TPU) 互转 to(device)

PyTorch .to(device) ↔ jax.devices()，数据搬运差异。

混合精度训练模型互转 amp / jnp.float16

PyTorch amp ↔ JAX jnp.float16，参数精度转换。

分布式训练模型互转 DDP / pmap

PyTorch DDP ↔ JAX pmap/shard_map，模型分片策略对应。

预训练模型 (BERT/GPT) 互转实战 HuggingFace

导出权重并在JAX中重建。

Vision Transformer (ViT) 互转实战 Flax ViT

PyTorch ViT ↔ Flax ViT 完整转换流程。

自定义层与复杂结构互转动态图特性

处理自定义forward、条件分支、循环等。

模型转换后的微调与部署生产格式

加载转换模型，微调训练并导出。

常见错误与调试指南 Key/Shape/NaN

Key不匹配、shape错误、NaN/Inf、随机种子不一致等排查修复。