做深度学习第一步不是搭模型而是选框架。PyTorch、TensorFlow、JAX 这三套 API 设计思路差异非常大代码迁移成本高选错了后面写训练循环、做部署都要返工。这篇文章以“7 个框架全景”为背景重点把三个主力框架的核心 API 拆开对比环境安装、张量操作、自动微分、模型构建、训练循环、数据加载、部署导出、资源占用和问题排查全部用最小可运行示例验证。如果你正打算从 TensorFlow 迁到 PyTorch或者想试 JAX 的函数式变换但一直没上手这篇可以直接收藏。先给结论没有“最强框架”只有“最匹配场景的 API”。PyTorch 适合研究和快速迭代TensorFlow 适合生产系统和端侧部署JAX 适合需要自动微分、自动向量化、自动并行化的高性能数值计算。下面用一套“从环境到推理”的最小测试流程把三个框架的差异点逐个讲清楚。1. 7 个深度学习框架核心能力速览框架来源/社区核心 API 风格主要定位PyTorchMeta 发起现由 PyTorch Foundation 管理动态图nn.Module autograd科研、快速原型、工业推理TensorFlowGoogle动态图 Keras 高层 API生产系统、端侧部署、大规模分布式JAXGoogle函数式变换jnp grad/jit/vmap/pmap高性能数值计算、科研算法复现KerasGoogle / 社区高层 API支持多后端快速搭建、教学、迁移到不同后端PaddlePaddle百度动态图为主动静统一工业应用与科研中文生态MindSpore华为动静统一自动并行昇腾硬件生态、企业 AIMXNetApacheGluon 动态图接口历史使用广泛当前社区活跃度明显下降这 7 个框架里PyTorch、TensorFlow、JAX 的开源社区最活跃也是国内外论文复现和工程落地的绝对主流。Keras 现在更像一个“前端接口层”可以跑在 TensorFlow、JAX 甚至 PyTorch 后端上。PaddlePaddle、MindSpore 在特定硬件和中文工业场景里有很强支持。MXNet 今天主要用于维护老项目新项目不建议再选。2. 框架选型什么时候用哪个选框架不要看热度要看你要做什么。PyTorch 最值得选的理由是“调试直接”。模型就是一个普通 Python 对象前向传播是一段普通 Python 代码print、breakpoint、pdb 都能直接用。论文复现、快速验证新想法、做多轮实验PyTorch 的效率最高。PyTorch 的模型结构定义和动态控制流非常自然RNN、Transformer、扩散模型这类结构写起来都不费劲。TensorFlow 的强项在“生产链路完整”。从 TF Serving、TFLite、TensorFlow.js 到 TFX 流水线训练到部署的工程件齐全。Keras 高层接口让模型搭建非常快适合标准化团队协作和产品化落地。如果团队已经有完整的 Kubernetes 和模型服务基础设施TensorFlow 仍然是可靠选择。JAX 要接受的是一套完全不同的思维没有“模型对象”一切是纯函数加参数数组没有model.fit()训练循环要自己写换来的好处是grad、jit、vmap、pmap这种组合式变换在强化学习、分子动力学、贝叶斯建模、大模型分布式并行等场景里效率极高。Google DeepMind 的很多研究项目和开源库都基于 JAX。使用边界也要说清楚训练数据必须来源合法模型权重要看开源许可证涉及人脸、声音、隐私数据时要确认授权部署 API 服务要限制访问范围任何情况下都不要用框架去绕过安全限制或做侵权内容生成。3. 环境准备与安装三套框架都要求先确认 Python、CUDA、cuDNN 版本匹配。装不上 GPU 版通常是 CUDA 版本不一致或者驱动太旧排查顺序固定为驱动 - CUDA - Python - 框架。通用检查命令python --version nvidia-smi nvcc --version驱动只要满足 CUDA 版本即可。注意nvidia-smi显示的是驱动支持的 CUDA 版本不一定是本机安装的 CUDA toolkit 版本两个概念不要混淆。PyTorch 安装推荐用官方命令生成器# 不要直接复制去 pytorch.org 根据系统、CUDA 版本生成命令 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121TensorFlow 2.18 的 Linux GPU 安装推荐元包方式# Linux GPU 版CUDA 相关依赖由元包统一管理 pip install tensorflow[and-cuda] # 仅 CPU 版 pip install tensorflow-cpuJAX 的 GPU 安装要区分 CUDA 版本# CPU 版 pip install -U jax # CUDA 12 版 pip install -U jax[cuda12] # CUDA 11 版 pip install -U jax[cuda11]安装完不要急着写模型先做硬件验证# PyTorch import torch print(PyTorch, torch.__version__) print(CUDA available:, torch.cuda.is_available()) print(GPU:, torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU) # TensorFlow import tensorflow as tf print(TensorFlow, tf.__version__) print(GPU:, tf.config.list_physical_devices(GPU)) # JAX import jax print(JAX, jax.__version__) print(Devices:, jax.devices())JAX 如果输出CpuDevice说明 GPU 驱动或 CUDA 库没配对。TensorFlow 2.18 如果只显示 CPU重点检查tensorflow[and-cuda]是否安装成功而不是只装了tensorflow基础包。4. 核心张量 API 对比三者的核心张量类型分别是torch.Tensor、tf.Tensor、jax.ArrayJAX 统一用jnp.array创建。API 设计差异在创建、转换、设备管理上非常明显。import torch import tensorflow as tf import jax.numpy as jnp # PyTorch pt torch.tensor([1.0, 2.0, 3.0]) print(pt.dtype, pt.device) # TensorFlow tf_t tf.constant([1.0, 2.0, 3.0]) print(tf_t.dtype, tf_t.device) # JAX jx jnp.array([1.0, 2.0, 3.0]) print(jx.dtype, jx.device())三者的设计差异PyTorch 的张量是“可变的”。你可以原地改数据、移动设备但自动微分需要参与求导的张量显式设置requires_gradTrue默认关闭。TensorFlow 的张量默认“不可变”但tf.Variable是可变对象。Keras 模型里的权重就是tf.Variable。JAX 的数组默认“不可变”。每次运算返回新数组用jax.numpy替代 NumPy但 API 和 NumPy 高度一致迁移成本最低。张量形状和类型转换# PyTorch 查看和转换 pt torch.randn(4, 8) print(pt.shape, pt.size()) pt_np pt.numpy() # 注意 requires_gradTrue 时不能直接转换 pt2 torch.from_numpy(pt_np) # TensorFlow 查看和转换 print(tf_t.shape) tf_np tf_t.numpy() # JAX 查看和转换 print(jx.shape) jx_np jnp.asarray(jx)设备控制是三框架 API 差异最大的地方之一。PyTorch 使用.to(cuda)显式搬运张量TensorFlow 在strategy.scope()或 Keras 里自动处理JAX 更彻底数据默认就在加速器上普通jnp运算会自动选择可用设备。PyTorch 新手最常见的 Bug 就是把 CPU 张量直接传给 CUDA 模型报 mismatch 错误。建议统一写成model model.to(device)并把输入输出的设备逻辑封装到训练函数里。5. 自动微分 API 对比自动微分是深度学习框架的核心三者的实现思路完全不同。PyTorch 采用“动态计算图 反向传播”。张量开启requires_gradTrue后前向执行时自动记录梯度函数调用backward()后梯度回传。代码看起来和普通数值计算一致这是它容易上手的关键。import torch x torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) y (x ** 2).sum() y.backward() print(PyTorch grad:, x.grad) # [2.0, 4.0, 6.0]TensorFlow 用tf.GradientTape的上下文管理器显式记录前向计算。import tensorflow as tf x tf.Variable([1.0, 2.0, 3.0]) with tf.GradientTape() as tape: y tf.reduce_sum(x ** 2) grad tape.gradient(y, x) print(TensorFlow grad:, grad.numpy())JAX 用纯函数变换jax.grad没有动态图也没有“反向传播”这个动作而是直接对损失函数求梯度。要就求一阶写jax.grad要求“损失值和梯度一起拿”用jax.value_and_grad。import jax import jax.numpy as jnp def loss_func(x): return jnp.sum(x ** 2) x jnp.array([1.0, 2.0, 3.0]) grad jax.grad(loss_func)(x) print(JAX grad:, grad)JAX 的jax.jit编译、jax.vmap向量化、jax.pmap多设备并行都是“函数变换”和 Python 原来的控制流不是一回事。写 JAX 时要避免用纯 Python 的if/for处理张量分支尽量用jnp.where、jax.lax.scan这类可变换结构否则编译时机和性能都会有坑。6. 模型构建 API 对比PyTorch 用nn.Module。模型是类forward方法定义前向子模块自动收集参数。import torch.nn as nn class MLP(nn.Module): def __init__(self, in_dim, hidden_dim, out_dim): super().__init__() self.fc1 nn.Linear(in_dim, hidden_dim) self.relu nn.ReLU() self.fc2 nn.Linear(hidden_dim, out_dim) def forward(self, x): return self.fc2(self.relu(self.fc1(x)))TensorFlow 高层接口是 Keras。Sequential适合顺序结构Model适合多输入多输出和自定义结构。import tensorflow as tf model tf.keras.Sequential([ tf.keras.layers.Dense(hidden_dim, activationrelu), tf.keras.layers.Dense(out_dim) ])TensorFlow 还可以继承tf.keras.Model写自定义层和自定义前向风格上和 PyTorch 的nn.Module接近。团队内部要想清楚“统一走 Keras 高层接口”还是“自定义代码”两种混用会让维护成本上升。JAX 本身没有内置模型类生态里最常用的是 Flax。Flax 用nn.Module但核心是“参数初始化函数 apply 方法”模型实例只是配置描述不保存权重。import flax.linen as nn import jax.numpy as jnp class MLP(nn.Module): hidden_dim: int out_dim: int nn.compact def __call__(self, x): x nn.Dense(self.hidden_dim)(x) x nn.relu(x) x nn.Dense(self.out_dim)(x) return x model MLP(hidden_dim64, out_dim1) params model.init(jax.random.PRNGKey(0), jnp.ones((1, 10))) pred model.apply(params, jnp.ones((1, 10)))Flax 的params是一个独立字典训练时通过参数传递更新。这种设计一开始会不习惯但配合optax做参数更新时非常清晰。JAX 生态里也有 Haiku、Equinox 等替代库选型前先看团队共识。7. 训练循环 API 对比这是三个框架差异最明显、也是迁移成本最高的部分。PyTorch 的训练循环完全手写逻辑全在你控制之下optimizer torch.optim.Adam(model.parameters(), lr1e-3) loss_fn torch.nn.MSELoss() for epoch in range(num_epochs): for x_batch, y_batch in dataloader: optimizer.zero_grad() pred model(x_batch) loss loss_fn(pred, y_batch) loss.backward() optimizer.step()TensorFlow 有两种训练方式。不想写轮子就用model.fitmodel.compile(optimizeradam, lossmse) model.fit(x_train, y_train, epochs10, batch_size32, validation_split0.1)需要精细控制梯度时用GradientTape自定义训练循环optimizer tf.keras.optimizers.Adam(learning_rate1e-3) loss_fn tf.keras.losses.MeanSquaredError() for epoch in range(num_epochs): for x_batch, y_batch in dataset: with tf.GradientTape() as tape: pred model(x_batch, trainingTrue) loss loss_fn(y_batch, pred) grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables))JAX 是“自己写一切”但框架提供了可组合的变换。下面是一个最小训练步配合optaximport optax import jax def loss_fn(params, x_batch, y_batch): pred model.apply(params, x_batch) pred pred.reshape(y_batch.shape) return jnp.mean((pred - y_batch) ** 2) optimizer optax.adam(learning_rate1e-3) opt_state optimizer.init(params) jax.jit def train_step(params, opt_state, x_batch, y_batch): loss, grads jax.value_and_grad(loss_fn)(params, x_batch, y_batch) updates, opt_state optimizer.update(grads, opt_state, params) params optax.apply_updates(params, updates) return params, opt_state, loss注意jax.jit装饰后传入的数据必须是数组而不是 Dataset 迭代器所以 JAX 的数据加载通常先取“一块 numpy/tf.data 数据”再交给编译后的train_step。这也是 JAX 和 PyTorch 训练流程差异最大的地方。三者的选择标准可以这样记PyTorch 保留最大控制权且调试直接TensorFlow 的fit生产集成方便但自定义逻辑需要绕一下JAX 追求函数式纯变换适合手写科研算法但初始学习成本最高。8. 数据加载与预处理 API 对比数据管道在三大框架中各自独立接口不通用。PyTorch 用DatasetDataLoader。自定义 Dataset 只需实现__len__和__getitem__DataLoader 自动负责 batch、打乱、多进程加载。from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, x, y): self.x x self.y y def __len__(self): return len(self.x) def __getitem__(self, idx): return self.x[idx], self.y[idx] dataloader DataLoader(MyDataset(x_train, y_train), batch_size32, shuffleTrue, num_workers4)TensorFlow 用tf.data.Dataset。它的优势是自带管道优化prefetch、map、batch、cache都可以链式调用还能配合TFRecord做大规模数据流。import tensorflow as tf dataset tf.data.Dataset.from_tensor_slices((x_train, y_train)) dataset dataset.shuffle(buffer_size1000) dataset dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)JAX 没有专用的 DataLoader社区常用tf.data或grainDeepMind 开源的数据加载库。常见做法是先用tf.data.Dataset完成 map/batch/prefetch再用ds.as_numpy_iterator()喂给 JAX 训练循环。注意 JAX 训练循环通常用for batch in dataset:但batch是 numpy 数组直接传给jitted函数即可。import tensorflow as tf dataset tf.data.Dataset.from_tensor_slices((x_train, y_train)) dataset dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE) for x_batch, y_batch in dataset.as_numpy_iterator(): params, opt_state, loss train_step(params, opt_state, x_batch, y_batch)从数据加载 API 来看PyTorch 灵活但多进程配置要调试TensorFlow 工程化强但 API 层级较多JAX 没有标准答案靠组合。9. 部署与生态接口对比训练结束后部署路径决定了框架选型是否成功。PyTorch 常用导出方式是 TorchScript 和 ONNX。TorchScript 把模型编译为可序列化图适合 C 调用ONNX 是把模型迁移到其他运行时的重要通道很多加速卡厂商都支持 ONNX 导入。# PyTorch 导出 ONNX model.eval() dummy_input torch.randn(1, 10) torch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output])TensorFlow 部署是它的传统强项。SavedModel 是标准格式配合 TF Serving 直接提供 gRPC/HTTP 推理服务TFLite 适合移动端、嵌入式TensorFlow.js 能跑在浏览器和 Node.js。模型转换路径清晰从训练到生产不用切换体系。# TensorFlow 导出 SavedModel model.export(saved_model_dir)JAX 部署路径相对“年轻”。常用方案是jax2tf把 JAX 函数转成 TensorFlow 计算图再做 SavedModel 导出也有团队直接在生产环境用 XLA 编译的 JAX 函数做推理服务。JAX 在分布式并行推理上能力很强但推理基础设施需要自己搭没有 TensorFlow Serving 那种开箱即用的组件。如果你打算把模型接到 API 平台PyTorch 和 TensorFlow 都可以先用 ONNX/SavedModel 转换再交给专用推理服务。JAX 则要提前验证目标推理平台是否支持 XLA 或jax2tf转换否则部署环节会卡住。10. 资源占用与性能观察框架本身不会直接告诉你“显存够不够”要自己观察和分析。通用观察工具watch -n 1 nvidia-smiPyTorch 还可以在代码里打印显存分配print(fallocated: {torch.cuda.memory_allocated() / 1024**2:.1f} MB) print(freserved: {torch.cuda.memory_reserved() / 1024**2:.1f} MB)TensorFlow 默认会预占大量显存调试时可以改成按需增长gpus tf.config.list_physical_devices(GPU) if gpus: tf.config.experimental.set_memory_growth(gpus[0], True)JAX 查看设备数量import jax print(jax.device_count())显存占用主要受四个因素影响batch size、输入尺寸分辨率/序列长度、模型参数量、优化器状态。增大 batch 是训练速度收益最明显的手段但显存压力会同步上涨。如果显存不足优先减 batch size而不是降分辨率。减 batch 还不行再用梯度累加模拟大 batch。混合精度是另一个常用手段PyTorch 用torch.cuda.ampTensorFlow 用mixed_float16JAX 配合jax.disable_float32()或显式使用bfloat16。CPU 和 GPU 的性能差距很难给统一数字因为运算类型、数据量、环境都不同。要观察就固定数据规模分别跑 20 个 batch记录耗时和显存变化。对比时用同一套超参不要一边带编译优化一边不带否则对比结果没有参考价值。11. 常见问题与排查方法问题现象可能原因排查方式解决方案PyTorch 装完torch.cuda.is_available()为 FalseCUDA 版本与 PyTorch wheel 不匹配nvidia-smi看驱动检查安装命令 index-url去 pytorch.org 重新生成匹配 CUDA 的安装命令TensorFlow 2.18 装完检测不到 GPU缺少 CUDA/cuDNN 运行库检查是否安装tensorflow[and-cuda]Linux 安装tensorflow[and-cuda]Windows 对照官方文档配 CUDA DLLJAXjax.devices()只显示 CPUJAX CUDA 版未安装或库没配对打印jax.__version__确认安装的是jax[cuda12]等 GPU 包用对应 CUDA 版本的jax[cuda12]/jax[cuda11]重新安装训练时显存不足 OOMbatch size 过大、输入尺寸过大、优化器状态过多nvidia-smi看占用峰值减小 batch size使用梯度累加、混合精度、gradient checkpointingDataLoader 多进程卡死PyTorch Windows 下num_workers配置不当把num_workers调为 0 测试将启动代码放入if __name__ __main__:按系统调整num_workersPyTorch 加载老模型报错PyTorch 2.6 起torch.load默认weights_onlyTrue检查加载代码手动指定weights_onlyTrue或对可信权重使用完整加载并明确处理反序列化风险JAX 训练在 CPU/GPU 之间跳数据不是数组进入了不支持变换的 Python 控制流打印训练输入类型统一用jnp数组避免在jit装饰函数里用 Python 原生if/for判断张量model.fit效果正常自定义 GradientTape 报错变量没有用tf.Variable包装检查模型参数是否在trainable_variables自定义层和模型都继承tf.keras.Model让框架托管参数排查原则是“先环境后代码”。报错先看驱动、CUDA、Python 版本是否匹配再看数据形状和设备是否一致。三套框架的报错信息里都会给出设备、张量形状和具体操作位置不要只看第一行。12. 最佳实践与使用建议工程上不管是 PyTorch、TensorFlow 还是 JAX以下做法都适用。第一第一次跑新环境先小规模验证。不要直接上完整模型和大 batch先跑 1 个 batch、10 步训练确认前向、反向、优化器、保存全部能通再扩规模。这样能快速区分“环境问题”和“算法问题”。第二环境隔离和版本固定。用 conda 或 venv 为每个项目建立独立环境要求项目里记录 Python、框架、CUDA、关键依赖的精确版本。框架升级带来的兼容性问题比多数模型本身的问题更难排查。第三目录规范。建议按data/、models/、outputs/、src/分层管理原始数据、模型权重、日志、训练脚本分开。JAX 和 PyTorch 的模型权重格式不通用直接拆目录存params.pt、saved_model、params.pkl避免一个目录堆满二进制文件。第四批量训练任务要加日志和断点。PyTorch 训练循环里加torch.save断点非常自然TensorFlow 用ModelCheckpoint回调JAX 需要自己把params和opt_state序列化。没有断点机制就大规模训练任何一个节点中断都会浪费大量算力。第五模型保存和加载要适配版本。PyTorch 2.6 起torch.load默认weights_onlyTrue加载旧权重时先确认反序列化安全性。TensorFlow 的 SavedModel 和 Keras.h5格式不要混用JAX 生态的权重一般配合 Flax/optax 结构保存。最后合规提醒要前置。训练数据、人脸数据、语音数据、版权素材都要确认授权模型部署为 API 服务时要加访问控制对外不能无鉴权裸奔使用开源模型权重先检查许可证。做技术验证没问题公开上线或商用前必须走法务和合规检查。13. 总结与下一步这篇文章的核心结论是三条路对应三种思维方式PyTorch 让你像写普通 Python 一样写模型和训练逻辑调试成本最低TensorFlow 给你从训练到部署的最完整工程链路适合团队标准化交付JAX 让你用函数变换组合出高效计算流程适合科研算法和需要极致并行控制的项目。三者的核心 API 差异集中在张量可变性、自动微分方式、模型组织形式、训练循环控制权和数据加载方式五个维度。建议你拿到任何新框架都先跑一遍同样的最小测试安装验证、张量创建、求梯度、搭一个两层的 MLP、写一个 10 步训练循环、导出模型。谁都能用这套测试在半小时内跑通整个链路。最容易踩的坑是“用 PyTorch 的思维写 JAX”或者“用 TensorFlow 的高层接口写完却想在自定义训练循环里接管所有参数”。框架之间的迁移不是改 API 名而是改代码的组织方式。下一步可以往三个方向扩展一是对比三者的分布式训练接口DistributedDataParallel、tf.distribute.Strategy、jax.pmap完全是三种抽象二是研究 ONNX 作为跨框架交换格式把 PyTorch 和 TensorFlow 模型统一部署三是深入 JAX 的vmap/jit组合它的性能上限和调试复杂度都值得单独写一篇。建议先把最小测试在三个框架上都跑通后续再按任务类型选型别急着在大项目里做一次性迁移。 PyTorch、TensorFlow、JAX三大深度学习框架核心API全对比
做深度学习第一步不是搭模型而是选框架。PyTorch、TensorFlow、JAX 这三套 API 设计思路差异非常大代码迁移成本高选错了后面写训练循环、做部署都要返工。这篇文章以“7 个框架全景”为背景重点把三个主力框架的核心 API 拆开对比环境安装、张量操作、自动微分、模型构建、训练循环、数据加载、部署导出、资源占用和问题排查全部用最小可运行示例验证。如果你正打算从 TensorFlow 迁到 PyTorch或者想试 JAX 的函数式变换但一直没上手这篇可以直接收藏。先给结论没有“最强框架”只有“最匹配场景的 API”。PyTorch 适合研究和快速迭代TensorFlow 适合生产系统和端侧部署JAX 适合需要自动微分、自动向量化、自动并行化的高性能数值计算。下面用一套“从环境到推理”的最小测试流程把三个框架的差异点逐个讲清楚。1. 7 个深度学习框架核心能力速览框架来源/社区核心 API 风格主要定位PyTorchMeta 发起现由 PyTorch Foundation 管理动态图nn.Module autograd科研、快速原型、工业推理TensorFlowGoogle动态图 Keras 高层 API生产系统、端侧部署、大规模分布式JAXGoogle函数式变换jnp grad/jit/vmap/pmap高性能数值计算、科研算法复现KerasGoogle / 社区高层 API支持多后端快速搭建、教学、迁移到不同后端PaddlePaddle百度动态图为主动静统一工业应用与科研中文生态MindSpore华为动静统一自动并行昇腾硬件生态、企业 AIMXNetApacheGluon 动态图接口历史使用广泛当前社区活跃度明显下降这 7 个框架里PyTorch、TensorFlow、JAX 的开源社区最活跃也是国内外论文复现和工程落地的绝对主流。Keras 现在更像一个“前端接口层”可以跑在 TensorFlow、JAX 甚至 PyTorch 后端上。PaddlePaddle、MindSpore 在特定硬件和中文工业场景里有很强支持。MXNet 今天主要用于维护老项目新项目不建议再选。2. 框架选型什么时候用哪个选框架不要看热度要看你要做什么。PyTorch 最值得选的理由是“调试直接”。模型就是一个普通 Python 对象前向传播是一段普通 Python 代码print、breakpoint、pdb 都能直接用。论文复现、快速验证新想法、做多轮实验PyTorch 的效率最高。PyTorch 的模型结构定义和动态控制流非常自然RNN、Transformer、扩散模型这类结构写起来都不费劲。TensorFlow 的强项在“生产链路完整”。从 TF Serving、TFLite、TensorFlow.js 到 TFX 流水线训练到部署的工程件齐全。Keras 高层接口让模型搭建非常快适合标准化团队协作和产品化落地。如果团队已经有完整的 Kubernetes 和模型服务基础设施TensorFlow 仍然是可靠选择。JAX 要接受的是一套完全不同的思维没有“模型对象”一切是纯函数加参数数组没有model.fit()训练循环要自己写换来的好处是grad、jit、vmap、pmap这种组合式变换在强化学习、分子动力学、贝叶斯建模、大模型分布式并行等场景里效率极高。Google DeepMind 的很多研究项目和开源库都基于 JAX。使用边界也要说清楚训练数据必须来源合法模型权重要看开源许可证涉及人脸、声音、隐私数据时要确认授权部署 API 服务要限制访问范围任何情况下都不要用框架去绕过安全限制或做侵权内容生成。3. 环境准备与安装三套框架都要求先确认 Python、CUDA、cuDNN 版本匹配。装不上 GPU 版通常是 CUDA 版本不一致或者驱动太旧排查顺序固定为驱动 - CUDA - Python - 框架。通用检查命令python --version nvidia-smi nvcc --version驱动只要满足 CUDA 版本即可。注意nvidia-smi显示的是驱动支持的 CUDA 版本不一定是本机安装的 CUDA toolkit 版本两个概念不要混淆。PyTorch 安装推荐用官方命令生成器# 不要直接复制去 pytorch.org 根据系统、CUDA 版本生成命令 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121TensorFlow 2.18 的 Linux GPU 安装推荐元包方式# Linux GPU 版CUDA 相关依赖由元包统一管理 pip install tensorflow[and-cuda] # 仅 CPU 版 pip install tensorflow-cpuJAX 的 GPU 安装要区分 CUDA 版本# CPU 版 pip install -U jax # CUDA 12 版 pip install -U jax[cuda12] # CUDA 11 版 pip install -U jax[cuda11]安装完不要急着写模型先做硬件验证# PyTorch import torch print(PyTorch, torch.__version__) print(CUDA available:, torch.cuda.is_available()) print(GPU:, torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU) # TensorFlow import tensorflow as tf print(TensorFlow, tf.__version__) print(GPU:, tf.config.list_physical_devices(GPU)) # JAX import jax print(JAX, jax.__version__) print(Devices:, jax.devices())JAX 如果输出CpuDevice说明 GPU 驱动或 CUDA 库没配对。TensorFlow 2.18 如果只显示 CPU重点检查tensorflow[and-cuda]是否安装成功而不是只装了tensorflow基础包。4. 核心张量 API 对比三者的核心张量类型分别是torch.Tensor、tf.Tensor、jax.ArrayJAX 统一用jnp.array创建。API 设计差异在创建、转换、设备管理上非常明显。import torch import tensorflow as tf import jax.numpy as jnp # PyTorch pt torch.tensor([1.0, 2.0, 3.0]) print(pt.dtype, pt.device) # TensorFlow tf_t tf.constant([1.0, 2.0, 3.0]) print(tf_t.dtype, tf_t.device) # JAX jx jnp.array([1.0, 2.0, 3.0]) print(jx.dtype, jx.device())三者的设计差异PyTorch 的张量是“可变的”。你可以原地改数据、移动设备但自动微分需要参与求导的张量显式设置requires_gradTrue默认关闭。TensorFlow 的张量默认“不可变”但tf.Variable是可变对象。Keras 模型里的权重就是tf.Variable。JAX 的数组默认“不可变”。每次运算返回新数组用jax.numpy替代 NumPy但 API 和 NumPy 高度一致迁移成本最低。张量形状和类型转换# PyTorch 查看和转换 pt torch.randn(4, 8) print(pt.shape, pt.size()) pt_np pt.numpy() # 注意 requires_gradTrue 时不能直接转换 pt2 torch.from_numpy(pt_np) # TensorFlow 查看和转换 print(tf_t.shape) tf_np tf_t.numpy() # JAX 查看和转换 print(jx.shape) jx_np jnp.asarray(jx)设备控制是三框架 API 差异最大的地方之一。PyTorch 使用.to(cuda)显式搬运张量TensorFlow 在strategy.scope()或 Keras 里自动处理JAX 更彻底数据默认就在加速器上普通jnp运算会自动选择可用设备。PyTorch 新手最常见的 Bug 就是把 CPU 张量直接传给 CUDA 模型报 mismatch 错误。建议统一写成model model.to(device)并把输入输出的设备逻辑封装到训练函数里。5. 自动微分 API 对比自动微分是深度学习框架的核心三者的实现思路完全不同。PyTorch 采用“动态计算图 反向传播”。张量开启requires_gradTrue后前向执行时自动记录梯度函数调用backward()后梯度回传。代码看起来和普通数值计算一致这是它容易上手的关键。import torch x torch.tensor([1.0, 2.0, 3.0], requires_gradTrue) y (x ** 2).sum() y.backward() print(PyTorch grad:, x.grad) # [2.0, 4.0, 6.0]TensorFlow 用tf.GradientTape的上下文管理器显式记录前向计算。import tensorflow as tf x tf.Variable([1.0, 2.0, 3.0]) with tf.GradientTape() as tape: y tf.reduce_sum(x ** 2) grad tape.gradient(y, x) print(TensorFlow grad:, grad.numpy())JAX 用纯函数变换jax.grad没有动态图也没有“反向传播”这个动作而是直接对损失函数求梯度。要就求一阶写jax.grad要求“损失值和梯度一起拿”用jax.value_and_grad。import jax import jax.numpy as jnp def loss_func(x): return jnp.sum(x ** 2) x jnp.array([1.0, 2.0, 3.0]) grad jax.grad(loss_func)(x) print(JAX grad:, grad)JAX 的jax.jit编译、jax.vmap向量化、jax.pmap多设备并行都是“函数变换”和 Python 原来的控制流不是一回事。写 JAX 时要避免用纯 Python 的if/for处理张量分支尽量用jnp.where、jax.lax.scan这类可变换结构否则编译时机和性能都会有坑。6. 模型构建 API 对比PyTorch 用nn.Module。模型是类forward方法定义前向子模块自动收集参数。import torch.nn as nn class MLP(nn.Module): def __init__(self, in_dim, hidden_dim, out_dim): super().__init__() self.fc1 nn.Linear(in_dim, hidden_dim) self.relu nn.ReLU() self.fc2 nn.Linear(hidden_dim, out_dim) def forward(self, x): return self.fc2(self.relu(self.fc1(x)))TensorFlow 高层接口是 Keras。Sequential适合顺序结构Model适合多输入多输出和自定义结构。import tensorflow as tf model tf.keras.Sequential([ tf.keras.layers.Dense(hidden_dim, activationrelu), tf.keras.layers.Dense(out_dim) ])TensorFlow 还可以继承tf.keras.Model写自定义层和自定义前向风格上和 PyTorch 的nn.Module接近。团队内部要想清楚“统一走 Keras 高层接口”还是“自定义代码”两种混用会让维护成本上升。JAX 本身没有内置模型类生态里最常用的是 Flax。Flax 用nn.Module但核心是“参数初始化函数 apply 方法”模型实例只是配置描述不保存权重。import flax.linen as nn import jax.numpy as jnp class MLP(nn.Module): hidden_dim: int out_dim: int nn.compact def __call__(self, x): x nn.Dense(self.hidden_dim)(x) x nn.relu(x) x nn.Dense(self.out_dim)(x) return x model MLP(hidden_dim64, out_dim1) params model.init(jax.random.PRNGKey(0), jnp.ones((1, 10))) pred model.apply(params, jnp.ones((1, 10)))Flax 的params是一个独立字典训练时通过参数传递更新。这种设计一开始会不习惯但配合optax做参数更新时非常清晰。JAX 生态里也有 Haiku、Equinox 等替代库选型前先看团队共识。7. 训练循环 API 对比这是三个框架差异最明显、也是迁移成本最高的部分。PyTorch 的训练循环完全手写逻辑全在你控制之下optimizer torch.optim.Adam(model.parameters(), lr1e-3) loss_fn torch.nn.MSELoss() for epoch in range(num_epochs): for x_batch, y_batch in dataloader: optimizer.zero_grad() pred model(x_batch) loss loss_fn(pred, y_batch) loss.backward() optimizer.step()TensorFlow 有两种训练方式。不想写轮子就用model.fitmodel.compile(optimizeradam, lossmse) model.fit(x_train, y_train, epochs10, batch_size32, validation_split0.1)需要精细控制梯度时用GradientTape自定义训练循环optimizer tf.keras.optimizers.Adam(learning_rate1e-3) loss_fn tf.keras.losses.MeanSquaredError() for epoch in range(num_epochs): for x_batch, y_batch in dataset: with tf.GradientTape() as tape: pred model(x_batch, trainingTrue) loss loss_fn(y_batch, pred) grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables))JAX 是“自己写一切”但框架提供了可组合的变换。下面是一个最小训练步配合optaximport optax import jax def loss_fn(params, x_batch, y_batch): pred model.apply(params, x_batch) pred pred.reshape(y_batch.shape) return jnp.mean((pred - y_batch) ** 2) optimizer optax.adam(learning_rate1e-3) opt_state optimizer.init(params) jax.jit def train_step(params, opt_state, x_batch, y_batch): loss, grads jax.value_and_grad(loss_fn)(params, x_batch, y_batch) updates, opt_state optimizer.update(grads, opt_state, params) params optax.apply_updates(params, updates) return params, opt_state, loss注意jax.jit装饰后传入的数据必须是数组而不是 Dataset 迭代器所以 JAX 的数据加载通常先取“一块 numpy/tf.data 数据”再交给编译后的train_step。这也是 JAX 和 PyTorch 训练流程差异最大的地方。三者的选择标准可以这样记PyTorch 保留最大控制权且调试直接TensorFlow 的fit生产集成方便但自定义逻辑需要绕一下JAX 追求函数式纯变换适合手写科研算法但初始学习成本最高。8. 数据加载与预处理 API 对比数据管道在三大框架中各自独立接口不通用。PyTorch 用DatasetDataLoader。自定义 Dataset 只需实现__len__和__getitem__DataLoader 自动负责 batch、打乱、多进程加载。from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self, x, y): self.x x self.y y def __len__(self): return len(self.x) def __getitem__(self, idx): return self.x[idx], self.y[idx] dataloader DataLoader(MyDataset(x_train, y_train), batch_size32, shuffleTrue, num_workers4)TensorFlow 用tf.data.Dataset。它的优势是自带管道优化prefetch、map、batch、cache都可以链式调用还能配合TFRecord做大规模数据流。import tensorflow as tf dataset tf.data.Dataset.from_tensor_slices((x_train, y_train)) dataset dataset.shuffle(buffer_size1000) dataset dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE)JAX 没有专用的 DataLoader社区常用tf.data或grainDeepMind 开源的数据加载库。常见做法是先用tf.data.Dataset完成 map/batch/prefetch再用ds.as_numpy_iterator()喂给 JAX 训练循环。注意 JAX 训练循环通常用for batch in dataset:但batch是 numpy 数组直接传给jitted函数即可。import tensorflow as tf dataset tf.data.Dataset.from_tensor_slices((x_train, y_train)) dataset dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE) for x_batch, y_batch in dataset.as_numpy_iterator(): params, opt_state, loss train_step(params, opt_state, x_batch, y_batch)从数据加载 API 来看PyTorch 灵活但多进程配置要调试TensorFlow 工程化强但 API 层级较多JAX 没有标准答案靠组合。9. 部署与生态接口对比训练结束后部署路径决定了框架选型是否成功。PyTorch 常用导出方式是 TorchScript 和 ONNX。TorchScript 把模型编译为可序列化图适合 C 调用ONNX 是把模型迁移到其他运行时的重要通道很多加速卡厂商都支持 ONNX 导入。# PyTorch 导出 ONNX model.eval() dummy_input torch.randn(1, 10) torch.onnx.export(model, dummy_input, model.onnx, input_names[input], output_names[output])TensorFlow 部署是它的传统强项。SavedModel 是标准格式配合 TF Serving 直接提供 gRPC/HTTP 推理服务TFLite 适合移动端、嵌入式TensorFlow.js 能跑在浏览器和 Node.js。模型转换路径清晰从训练到生产不用切换体系。# TensorFlow 导出 SavedModel model.export(saved_model_dir)JAX 部署路径相对“年轻”。常用方案是jax2tf把 JAX 函数转成 TensorFlow 计算图再做 SavedModel 导出也有团队直接在生产环境用 XLA 编译的 JAX 函数做推理服务。JAX 在分布式并行推理上能力很强但推理基础设施需要自己搭没有 TensorFlow Serving 那种开箱即用的组件。如果你打算把模型接到 API 平台PyTorch 和 TensorFlow 都可以先用 ONNX/SavedModel 转换再交给专用推理服务。JAX 则要提前验证目标推理平台是否支持 XLA 或jax2tf转换否则部署环节会卡住。10. 资源占用与性能观察框架本身不会直接告诉你“显存够不够”要自己观察和分析。通用观察工具watch -n 1 nvidia-smiPyTorch 还可以在代码里打印显存分配print(fallocated: {torch.cuda.memory_allocated() / 1024**2:.1f} MB) print(freserved: {torch.cuda.memory_reserved() / 1024**2:.1f} MB)TensorFlow 默认会预占大量显存调试时可以改成按需增长gpus tf.config.list_physical_devices(GPU) if gpus: tf.config.experimental.set_memory_growth(gpus[0], True)JAX 查看设备数量import jax print(jax.device_count())显存占用主要受四个因素影响batch size、输入尺寸分辨率/序列长度、模型参数量、优化器状态。增大 batch 是训练速度收益最明显的手段但显存压力会同步上涨。如果显存不足优先减 batch size而不是降分辨率。减 batch 还不行再用梯度累加模拟大 batch。混合精度是另一个常用手段PyTorch 用torch.cuda.ampTensorFlow 用mixed_float16JAX 配合jax.disable_float32()或显式使用bfloat16。CPU 和 GPU 的性能差距很难给统一数字因为运算类型、数据量、环境都不同。要观察就固定数据规模分别跑 20 个 batch记录耗时和显存变化。对比时用同一套超参不要一边带编译优化一边不带否则对比结果没有参考价值。11. 常见问题与排查方法问题现象可能原因排查方式解决方案PyTorch 装完torch.cuda.is_available()为 FalseCUDA 版本与 PyTorch wheel 不匹配nvidia-smi看驱动检查安装命令 index-url去 pytorch.org 重新生成匹配 CUDA 的安装命令TensorFlow 2.18 装完检测不到 GPU缺少 CUDA/cuDNN 运行库检查是否安装tensorflow[and-cuda]Linux 安装tensorflow[and-cuda]Windows 对照官方文档配 CUDA DLLJAXjax.devices()只显示 CPUJAX CUDA 版未安装或库没配对打印jax.__version__确认安装的是jax[cuda12]等 GPU 包用对应 CUDA 版本的jax[cuda12]/jax[cuda11]重新安装训练时显存不足 OOMbatch size 过大、输入尺寸过大、优化器状态过多nvidia-smi看占用峰值减小 batch size使用梯度累加、混合精度、gradient checkpointingDataLoader 多进程卡死PyTorch Windows 下num_workers配置不当把num_workers调为 0 测试将启动代码放入if __name__ __main__:按系统调整num_workersPyTorch 加载老模型报错PyTorch 2.6 起torch.load默认weights_onlyTrue检查加载代码手动指定weights_onlyTrue或对可信权重使用完整加载并明确处理反序列化风险JAX 训练在 CPU/GPU 之间跳数据不是数组进入了不支持变换的 Python 控制流打印训练输入类型统一用jnp数组避免在jit装饰函数里用 Python 原生if/for判断张量model.fit效果正常自定义 GradientTape 报错变量没有用tf.Variable包装检查模型参数是否在trainable_variables自定义层和模型都继承tf.keras.Model让框架托管参数排查原则是“先环境后代码”。报错先看驱动、CUDA、Python 版本是否匹配再看数据形状和设备是否一致。三套框架的报错信息里都会给出设备、张量形状和具体操作位置不要只看第一行。12. 最佳实践与使用建议工程上不管是 PyTorch、TensorFlow 还是 JAX以下做法都适用。第一第一次跑新环境先小规模验证。不要直接上完整模型和大 batch先跑 1 个 batch、10 步训练确认前向、反向、优化器、保存全部能通再扩规模。这样能快速区分“环境问题”和“算法问题”。第二环境隔离和版本固定。用 conda 或 venv 为每个项目建立独立环境要求项目里记录 Python、框架、CUDA、关键依赖的精确版本。框架升级带来的兼容性问题比多数模型本身的问题更难排查。第三目录规范。建议按data/、models/、outputs/、src/分层管理原始数据、模型权重、日志、训练脚本分开。JAX 和 PyTorch 的模型权重格式不通用直接拆目录存params.pt、saved_model、params.pkl避免一个目录堆满二进制文件。第四批量训练任务要加日志和断点。PyTorch 训练循环里加torch.save断点非常自然TensorFlow 用ModelCheckpoint回调JAX 需要自己把params和opt_state序列化。没有断点机制就大规模训练任何一个节点中断都会浪费大量算力。第五模型保存和加载要适配版本。PyTorch 2.6 起torch.load默认weights_onlyTrue加载旧权重时先确认反序列化安全性。TensorFlow 的 SavedModel 和 Keras.h5格式不要混用JAX 生态的权重一般配合 Flax/optax 结构保存。最后合规提醒要前置。训练数据、人脸数据、语音数据、版权素材都要确认授权模型部署为 API 服务时要加访问控制对外不能无鉴权裸奔使用开源模型权重先检查许可证。做技术验证没问题公开上线或商用前必须走法务和合规检查。13. 总结与下一步这篇文章的核心结论是三条路对应三种思维方式PyTorch 让你像写普通 Python 一样写模型和训练逻辑调试成本最低TensorFlow 给你从训练到部署的最完整工程链路适合团队标准化交付JAX 让你用函数变换组合出高效计算流程适合科研算法和需要极致并行控制的项目。三者的核心 API 差异集中在张量可变性、自动微分方式、模型组织形式、训练循环控制权和数据加载方式五个维度。建议你拿到任何新框架都先跑一遍同样的最小测试安装验证、张量创建、求梯度、搭一个两层的 MLP、写一个 10 步训练循环、导出模型。谁都能用这套测试在半小时内跑通整个链路。最容易踩的坑是“用 PyTorch 的思维写 JAX”或者“用 TensorFlow 的高层接口写完却想在自定义训练循环里接管所有参数”。框架之间的迁移不是改 API 名而是改代码的组织方式。下一步可以往三个方向扩展一是对比三者的分布式训练接口DistributedDataParallel、tf.distribute.Strategy、jax.pmap完全是三种抽象二是研究 ONNX 作为跨框架交换格式把 PyTorch 和 TensorFlow 模型统一部署三是深入 JAX 的vmap/jit组合它的性能上限和调试复杂度都值得单独写一篇。建议先把最小测试在三个框架上都跑通后续再按任务类型选型别急着在大项目里做一次性迁移。