深度学习框架选型与安装：TensorFlow、PyTorch、JAX 对比与搭建指南

选框架这事儿，说白了跟挑工具差不多。我个人习惯是：先看项目需求，再挑趁手的。别被朋友圈吹的“最强框架”带偏了。你想想看，一个做移动端部署的团队跑去用 JAX，那不是给自己找麻烦吗？

一、三大框架的定位差异

PyTorch 主打“研究员友好”。动态图、调试方便、论文复现快。我在做图像生成模型时，大部分时间都泡在 PyTorch 里。

TensorFlow 强在“生产部署”。TFX 管线、TFLite、TensorFlow Serving，一条龙服务。你如果要做工业级推理，TF 依然是稳妥选择。

JAX 是个狠角色。函数式、自动微分、XLA 编译，跑大规模并行计算时优势明显。不过学习曲线偏陡，社区资源也少一些。

维度	PyTorch	TensorFlow	JAX
计算图风格	动态图（eager）	静态图 + 动态图（2.x）	函数式 + JIT
调试体验	自然（Python 原生）	需开启 eager 模式	对新手不友好
部署生态	TorchScript / ONNX / TorchServe	TF Serving / TFLite / TF.js	Passthrough + XLA
社区活跃度	极高，论文标配	高，工业应用多	增长中，学术偏多

二、选型到底看什么？

• 团队技术栈：团队熟悉 PyTorch 就别硬转 TF，除非部署强需求。
• 部署目标：移动端 → TF Lite；云服务 → TF Serving / TorchServe。
• 调试效率：快速原型 → PyTorch；需要严格优化 → JAX。
• 算力资源：JAX 需要较新 GPU 和 JIT 适配，老卡可能遇到兼容问题。

三、PyTorch 安装实战

PyTorch 安装其实最省心。官方提供了按 CUDA 版本区分的命令。我一般用 conda 或 pip 虚拟环境，免得污染系统包。

# 根据 CUDA 版本选择，这里以 CUDA 11.8 为例
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

验证安装：

python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

如果打印 True，恭喜，GPU 可用。经常有人问我为什么 torch.cuda.is_available() 返回 False —— 多半是 CUDA 版本没对齐，或者装了 CPU 版。我曾经也踩过这个坑，折腾了半天才发现 pip 默认装了 cpuonly。

四、TensorFlow 安装实战

TensorFlow 2.x 的 pip 包已经整合了 GPU 支持。但版本号要特别小心。比如 tensorflow==2.10 是最后一个原生支持 Windows GPU 的版本，之后的版本需要用 WSL2。

# Linux / macOS 通用
pip install tensorflow

# 如果要指定 CUDA 版本，可以尝试 tf-nightly 或从源码编译
pip install tensorflow==2.10

验证：

python -c "import tensorflow as tf; print(tf.__version__); print(tf.config.list_physical_devices('GPU'))"

如果列表为空，检查一下 nvidia-smi 输出，确认驱动和 CUDA 版本。我以前遇到过 tf 2.12 对 CUDA 11.8 支持有 bug，降级到 2.11 就好了。社区里这种事不少见。

五、JAX 安装实战

JAX 安装要注意：jax 是纯 CPU 版，jaxlib 提供 GPU 支持。必须手动指定对应的 CUDA 版本。

# CPU 版
pip install jax jaxlib

# GPU 版（以 CUDA 11.8 为例）
pip install --upgrade jax jaxlib==0.4.13+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

验证：

python -c "import jax; print(jax.__version__); print(jax.devices())"

输出 [cuda(id=0)] 就对了。JAX 对 GPU 型号很敏感，如果遇到 segmentation fault，八成是 cuDNN 或 CUDA 不匹配。我的建议是：严格对照 JAX 官方版本兼容表（不输出链接，此处仅指概念）。

六、一张图看懂选型决策

七、框架对比总览

特性	PyTorch (2.x)	TensorFlow (2.x)	JAX (0.4.x)
计算图	动态图，支持 torch.compile	默认 eager，可开启 tf.function	函数式 + jit 编译
API 风格	面向对象，nn.Module	Keras 高层 / 底层 API	纯函数，无状态
分布式训练	torch.distributed / FSDP	tf.distribute.Strategy	pmap / shard_map
推理部署	TorchScript / ONNX / TensorRT	Serving / Lite / JS	XLA 编译 + 自定义推理
社区生态	HuggingFace 首选，论文复现多	工业生态完善	DeepMind 等研究机构使用

八、避坑指南

• 版本不匹配：我曾经在同一个服务器上同时装 PyTorch 1.13 和 2.0，结果 CUDA 库冲突。后来统一用 conda 环境解决了。
• CPU 版误装：很多人用 pip install tensorflow 后才发现是 CPU 版。建议加 --extra-index-url 指定 GPU 版。
• JAX 的 CUDA 版本对齐：JAX 对 CUDA 和 cuDNN 的小版本要求很严。装之前先去官方 release notes 查一下对照表。
• 多框架共存：如果必须同时用 PyTorch 和 TF，建议用 Docker 容器，或者使用 conda 虚拟环境，每个环境只装一个框架。