深度学习框架对比

框架之争本质是编程范式之争：PyTorch 用 imperative 动态图换迭代速度，TensorFlow 用图执行换部署确定性，JAX 用纯函数加编译变换换极致性能。2025 年前后三者哲学趋同（torch.compile、tf.function、共享 XLA 编译底座），选型最终落在生态与部署约束，而非 API 好恶。

选型速查

场景	推荐	关键理由
研究、论文复现、快速原型	PyTorch	define-by-run 调试直观，学界代码事实标准
LLM / 大模型训练	PyTorch	FSDP、Megatron、DeepSpeed 工具链成熟
TPU 大规模实验	JAX	pjit 与 TPU 协同设计，编译优化最激进
移动 / 边缘部署	TensorFlow（LiteRT）	运行时体积小、硬件 delegate 全；PyTorch 侧对应 ExecuTorch
存量工业系统、浏览器	TensorFlow	TF Serving、TFX、TF.js 集成完整
高性能科学计算	JAX	jit/vmap/grad 可组合变换

设计哲学：三种范式

PyTorch：define-by-run

代码即模型：每条 Python 语句执行时实时构建计算图，backward 后图即弃。控制流就是 Python 控制流，形状、分支、循环无需特殊语法。求导由 torch.autograd 引擎沿动态图回传完成。

import torch
 
model = torch.nn.Sequential(torch.nn.Linear(784, 256), torch.nn.ReLU(),
                            torch.nn.Linear(256, 10))
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
 
for x, y in loader:          # 每个 batch 重新执行 Python，图随运行构建
    loss = torch.nn.functional.cross_entropy(model(x), y)
    loss.backward()          # 沿动态图回传
    opt.step(); opt.zero_grad()

PyTorch 2.x 引入 torch.compile：TorchDynamo 捕获 FX 图、Inductor 生成 fused Triton kernel，动态图框架补上编译优化——imperative 体验与静态图性能从此可兼得。分层架构细节见 PyTorch 架构总览。

TensorFlow：从静态图折返 eager

TF1.x 时代先显式构建 Graph 再 session.run，性能与部署友好但调试痛苦。TF2 默认 eager 执行，tf.function 装饰器按需将函数编译回图，Keras 成为官方高层 API：

import tensorflow as tf
 
model = tf.keras.Sequential([
    tf.keras.layers.Dense(256, activation="relu"),
    tf.keras.layers.Dense(10),
])
model.compile(optimizer="adam", loss="sparse_categorical_crossentropy")
model.fit(train_ds, epochs=5)  # fit 内部经 tf.function 编译为图执行

代价是两次 API 大改留下的兼容包袱（tf.compat、session 时代的存量代码），以及 tf.function 内部 AutoGraph 的隐式追踪陷阱。

JAX：纯函数与编译变换

JAX 把数值程序当纯函数处理，框架能力体现为可组合的函数变换：jit（XLA 编译）、grad（求导）、vmap（自动批处理）、pjit/pmap（并行）。没有 Module、没有 optimizer 对象、没有全局状态：

import jax, optax
 
optimizer = optax.adam(1e-3)
 
def loss_fn(params, x, y):     # 纯函数：无副作用、无全局状态
    return optax.softmax_cross_entropy_with_integer_labels(
        forward(params, x), y).mean()
 
@jax.jit
def step(params, opt_state, x, y):
    grads = jax.grad(loss_fn)(params, x, y)   # grad 是函数变换，不是方法
    updates, opt_state = optimizer.update(grads, opt_state)
    return optax.apply_updates(params, updates), opt_state

随机数必须显式传递 PRNG key；模型与训练状态全部显式参数化。换来的是最激进的编译期优化（fusion、GSPMD 并行）。

范式对照

维度	PyTorch	TensorFlow 2	JAX
编程范式	imperative、面向对象	混合：eager + Keras 声明式	函数式、纯函数
图构建	运行时动态建图	tf.function 编译静态图	jit 整体编译（XLA）
求导	autograd 引擎	GradientTape 上下文	grad 函数变换
随机数	全局 RNG	全局 RNG	显式 PRNG key
批处理/并行	手写循环或库封装	少量自动机制	vmap/pjit 一等公民

调试与开发体验

PyTorch：pdb / IDE 断点直接可用，报错栈就是 Python 栈，print 即时可见。训练循环全透明，改一行立刻生效。
TensorFlow 2：eager 模式已可调试，但 tf.function 编译区内的报错信息晦涩，AutoGraph 对 Python 子集的追踪与直觉有偏差；历史 API 演进增加心智负担。
JAX：不加 jit 时是普通 NumPy 代码，可调试；一旦 jit，错误以编译期 shape mismatch 形式出现且定位困难。PRNG 显式传递与 purity 约束抬高了入门成本，但换来的可复现性是研究刚需。

生态版图

层次	PyTorch	TensorFlow	JAX
神经网络库	torch.nn 内建	Keras（3.x 起多后端）	Flax（Haiku 已归档）
预训练模型	HuggingFace Hub、torchvision	Keras Applications、TF Hub	Hub 上的 Flax 权重
训练封装	Lightning、Accelerate	Keras fit、TFX	Flax Train、机构内部库
数据管道	DataLoader	tf.data	Grain、复用 TFDS
分布式	分布式训练总览中 torchrun + FSDP 等工具链	tf.distribute	pjit、GSPMD

HuggingFace Transformers 以 PyTorch 为第一公民，JAX 与 TF 作为次级后端存在——这一事实深刻影响了三者的实际使用者规模。

推理与部署

PyTorch：生产 serving 常用 NVIDIA Triton（libtorch backend）；官方 TorchServe 维护放缓。跨框架部署走 ONNX 模型导出；边缘端有 ExecuTorch。LLM 推理生态（vLLM、SGLang）直接构建在 PyTorch 之上，权重分发则依赖 safetensors/GGUF 等格式（见推理模型格式）。
TensorFlow：部署链路最完整——TF Serving（gRPC/REST 工业标配）、LiteRT（原 TFLite，移动端 delegate 丰富）、TF.js（浏览器）、TFX（端到端 pipeline）。
JAX：无官方 serving 栈。jax.export 产出 StableHLO 交给 XLA 运行时，或经 ONNX 转换、HuggingFace 导出权重间接落地——部署依赖导出链路是选型时的主要顾虑。

社区趋势（2025 年前后）

学术界 PyTorch 绝对主导：顶会新论文的复现代码默认 PyTorch，研究生入门即学。「论文代码 = PyTorch 代码」的网络效应仍在自我强化。
JAX 影响力集中于研究机构：DeepMind 全线采用（AlphaFold 系列即 JAX 实现），Anthropic 等前沿实验室用于大模型训练。研究影响力远大于生产存量。
TensorFlow 工业存量深厚：Google 生产体系、移动端 LiteRT、大量存量企业模型仍运行于 TF Serving。新项目份额下滑，但离「消亡」很远。
哲学趋同：PyTorch 用 torch.compile 补编译，JAX 生态（Flax、Equinox）不断降低函数式门槛，Keras 3 反向支持 PyTorch/JAX/TF 三后端。范式之争的胜负已不重要，工程重心上移到编译器（XLA、Inductor/Triton）与系统层（vLLM、Megatron）。

参考链接

三大深度学习框架对比（知乎）

知识花园

探索