TensorFlow 和 PyTorch 的选型问题,几乎是每个入坑深度学习的人都会遇到的第一道坎。一个是 Google 在 2015 年开源的生产级框架,一个是 Meta AI 在 2016 年推出的动态图框架,经过多年迭代,两者都已经非常成熟,但“到底学哪个、用哪个”依然没有标准答案。这篇文章不打算站队,而是把两个框架的定位、安装方式、实战写法、部署生态和常见坑一次讲清楚,让你能根据自己手头的任务做决定。
先说结论方向:如果目标是科研、论文复现、快速验证新模型,PyTorch 在研究和教育社区的生态更友好;如果目标是把模型部署到生产系统、服务端推理、移动端和嵌入式设备,TensorFlow 的产业化链路更成熟。但这不是绝对的,近两年 PyTorch 在部署侧的进展也很快,TensorFlow 也吸收了 Keras、动态图等大量易用特性。后面会拆开细说。
本文会覆盖以下实操内容:搭建 Python 虚拟环境、安装 TensorFlow 和 PyTorch(CPU / GPU 两种方式)、用两个框架分别跑通一个简单的线性回归和一个 MNIST 手写数字识别、对比训练接口和部署流程、给出常见报错排查清单。全程使用可复制的代码块,你可以直接跟着跑。看完之后,你应该能判断自己当前的项目更适合哪一个框架。
1. 核心能力速览
先给一张速览表,把两个框架的关键信息放在一起对比。需要注意的是,两个框架迭代速度都很快,下面的参数和定位属于长期稳定的“骨架信息”,具体小版本能力要结合你实际安装的版本确认。
| 对比项 | TensorFlow | PyTorch |
|---|---|---|
| 开源方 | Meta AI | |
| 首个版本 | 2015 年 | 2016 年 |
| 默认图模式 | 动态执行(Eager Execution) | 动态计算图 |
| 高级 API | Keras(内置默认) | torch.nn + Lightning 等第三方 |
| 主要语言 | Python、C++ | Python、C++ |
| 官方部署工具 | TensorFlow Serving、TensorFlow Lite | TorchServe、ONNX Runtime、LibTorch |
| 移动端支持 | TFLite 生态成熟,跨 Android / iOS / MCU | PyTorch Mobile 可用,但生态不如 TFLite 丰富 |
| 模型可视化 | TensorBoard | TensorBoard(通过 torch.utils.tensorboard) |
| 分布式训练 | tf.distribute | torch.distributed |
| 典型用户 | 工业部署、端侧 AI、老牌生产系统 | 学术界、科研、AIGC、竞赛 |
| 学习曲线 | 接口较多,概念分层明显 | 接口更贴近 Python 习惯,上手较快 |
| 是否支持 CPU 训练 | 支持 | 支持 |
| 是否支持 GPU 训练 | 支持 NVIDIA GPU,部分框架版本支持 AMD ROCm | 支持 NVIDIA GPU,也有 ROCm 版本 |
| 是否是“默认选择” | 早期深度学习生产项目常选它 | 当前研究论文、开源模型默认权重大多用它 |
从表中可以看到,两个框架在基础能力上几乎没有绝对差距,差距主要体现在生态侧重和工程习惯上。
如果你看到这里还拿不定主意,可以先记住一个简单的判断标准:你要做的是“快速验证模型想法”,就从 PyTorch 开始;你要做的是“把模型稳定地跑在服务端或端侧”,就先研究 TensorFlow 的部署链路。当然,后端开发、平台工程师也经常会遇到“模型是 PyTorch 训练好的,但要转到 TensorFlow Serving 部署”的跨框架转换需求,这类工作我们在后面章节也会提到。
2. 适用场景与使用边界
2.1 哪些场景适合选 PyTorch
PyTorch 最大的优势在于“调试体验好”。因为默认是动态计算图,你可以像写普通 Python 代码一样打印中间张量、加断点、随时修改网络结构。这对于科研实验、课程作业、论文复现来说非常友好。
常见的适合选 PyTorch 的场景包括:
- 做学术研究、发论文、复现 GitHub 上的新模型。
- 做 NLP、CV、多模态、AIGC 相关实验,比如用 Transformer 做文本分类、用 Stable Diffusion 做图像生成。
- 参加 Kaggle 等数据竞赛,因为大多数 Top 方案和公开代码都是 PyTorch 写的。
- 需要频繁修改网络结构、自定义损失函数和训练循环的探索性项目。
2.2 哪些场景适合选 TensorFlow
TensorFlow 的优势在于“工程化链路完整”。从数据管道到训练、调参、版本管理、模型导出、上线推理,Google 都提供了配套工具。虽然没有早年那么“统治级”,但生产环境存量非常大。
常见的适合选 TensorFlow 的场景包括:
- 公司已有 TensorFlow 技术栈,需要和现有服务集成。
- 要做移动端、嵌入式、IoT 设备上的模型推理,TFLite 的兼容性更稳。
- 需要大规模分布式训练,并且希望靠框架内置能力减少自研成本。
- 需要长期维护的线上推理服务,TensorFlow Serving 的热加载、版本切换机制很成熟。
2.3 版权、隐私与安全使用边界
深度学习框架本身是开源工具,但使用时要注意几个边界:
- 训练数据来源要合法,尤其是人脸、声音、文本、图片等带版权或隐私属性的数据。
- 不要用模型生成或处理涉及他人肖像、声音的内容用于欺诈或误导。
- 使用第三方预训练模型时,检查许可证(如 Apache-2.0、MIT、CC-BY-NC 等),商用前确认是否允许。
- 本地训练和部署时,如果数据包含个人信息,建议做脱敏处理,并限制服务访问范围。
这里不展开法律条文,只强调一个原则:框架可以随便选,数据使用边界一定要清楚。
3. 环境准备与前置条件
3.1 通用检查清单
在安装 TensorFlow 或 PyTorch 之前,先确认自己的环境状态。下面的检查清单适用于 Windows、Linux 和 macOS,只是具体命令略有差异。
| 检查项 | 推荐要求 | 说明 |
|---|---|---|
| 操作系统 | Windows 10/11、Ubuntu 18.04 及以上、macOS 12+ | Linux 服务器最省心 |
| Python 版本 | 3.9 到 3.12 之间 | 不要直接上 3.13 之前的预览版,兼容性风险大 |
| GPU 驱动 | NVIDIA 驱动 470+ | 用nvidia-smi查看驱动版本和 CUDA 版本 |
| CUDA 工具包 | 11.8 或 12.x | 也可以不装系统级 CUDA,直接用 pip 的 CUDA 运行时 |
| 磁盘空间 | 至少 15GB 可用 | 框架本身不大,但 CUDA 库和模型数据占空间 |
| 内存 | 16GB 以上 | 8GB 也能跑小模型,但体验一般 |
| 代理/镜像 | 无强制要求 | 在国内建议用清华、阿里等 pip 镜像加速 |
3.2 用 conda 创建虚拟环境
强烈建议不要直接在系统 Python 环境里装深度学习框架,否则项目多了之后依赖会互相打架。推荐用 conda 或 venv 隔离环境。
# 安装 Miniconda 后执行 conda create -n deeplearn python=3.10 -y conda activate deeplearn创建好环境后,可以顺手升级 pip:
python -m pip install --upgrade pip如果 pip 下载慢,可以临时使用国内镜像:
pip config set global.index-url https://mirrors.aliyun.com/pypi/simple/3.3 确认 GPU 是否可用
Windows 和 Linux 下,打开终端执行:
nvidia-smi如果输出中出现了显卡型号、驱动版本和显存大小,说明 NVIDIA 驱动已经装好。之后安装框架时,只需要保证 pip 安装的 CUDA 版本和驱动兼容即可,不一定要手动装完整的 CUDA Toolkit。
在 macOS 上,只能做 CPU 训练,Apple Silicon 芯片可以走 MPS 后端。
4. TensorFlow 与 PyTorch 安装部署
4.1 安装 TensorFlow
TensorFlow 2.x 之后的安装非常统一,直接使用 pip 安装默认包即可。默认包包含了 GPU 支持的 CUDA 运行时,无需再单独安装tensorflow-gpu。
# CPU 版本 pip install tensorflow-cpu # GPU 版本(默认包,包含 CUDA 支持) pip install tensorflow如果你的机器是 NVIDIA GPU,建议直接执行pip install tensorflow。较新版本如 TensorFlow 2.18 对 CUDA 12 和 cuDNN 8.x 有对应兼容要求,具体以官方 release note 为准。
安装完成后,用一段 Python 代码验证:
import tensorflow as tf print("TensorFlow version:", tf.__version__) print("GPU available:", tf.config.list_physical_devices('GPU'))如果输出GPU available: []为空,说明 TensorFlow 没有检测到 GPU,需要检查驱动和 CUDA 运行时版本。如果只是跑 CPU,不影响后续学习。
4.2 安装 PyTorch
PyTorch 的安装命令建议从官网或者包管理器自动生成,不要死记硬背。基本格式是:
# CPU 版本 pip install torch torchvision torchaudio # CUDA 12.1 版本 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121在较新的 PyTorch 2.x 版本中,默认会安装包含 CUDA 支持的大包。如果你的网络访问官方 PyTorch 源比较慢,可以使用国内镜像源,但需要注意镜像源里的 CUDA 版本是否匹配你的驱动。
安装完成后验证:
import torch print("PyTorch version:", torch.__version__) print("CUDA available:", torch.cuda.is_available()) if torch.cuda.is_available(): print("GPU name:", torch.cuda.get_device_name(0))4.3 关于 PyTorch 2.6 的一个改动
如果你安装的是 PyTorch 2.6 及以后版本,需要注意官方做了一个默认行为变更:torch.load的weights_only参数默认值改为True。简单说,现在加载模型权重时更安全了,但如果你加载的是包含自定义类对象的旧存档,可能会报错。
解决办法是显式指定参数:
# 新的安全默认值 state_dict = torch.load("model.pth", weights_only=True) # 兼容旧存档 # state_dict = torch.load("model.pth", weights_only=False)这属于安全增强,在正常加载 PyTorch 官方权重时基本无感,但当你从网络下载一些老项目权重时,如果遇到WeightsUnpickler错误,优先排查是不是这个参数导致的。
5. 功能测试与效果验证
安装完成后,跑一个小例子是最快的验证方式。这里我们分别用 TensorFlow 和 PyTorch 实现一个简单的线性回归,再各跑一个 MNIST 手写数字识别小模型,用来对比两个框架的“手感”。
5.1 线性回归:TensorFlow 实现
import numpy as np import tensorflow as tf # 构造简单线性数据 x = np.linspace(-2, 2, 200).reshape(-1, 1).astype(np.float32) y = 3 * x + 1 + np.random.normal(0, 0.1, size=x.shape).astype(np.float32) # 定义模型 model = tf.keras.Sequential([ tf.keras.layers.Dense(1, input_shape=(1,)) ]) model.compile(optimizer='sgd', loss='mse') # 训练 history = model.fit(x, y, epochs=100, verbose=0) # 验证预测 print("预测值:", model.predict(np.array([[1.0]], dtype=np.float32)).item()) print("预期值: 约 4.0")这段代码把 200 个散点拟合成一条直线。如果训练正常,预测结果会接近 4.0。
5.2 线性回归:PyTorch 实现
import torch import torch.nn as nn # 构造数据 x = torch.linspace(-2, 2, 200).reshape(-1, 1) y = 3 * x + 1 + torch.randn(x.size()) * 0.1 # 定义模型 model = nn.Linear(1, 1) loss_fn = nn.MSELoss() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # 训练 for epoch in range(100): optimizer.zero_grad() pred = model(x) loss = loss_fn(pred, y) loss.backward() optimizer.step() # 验证预测 model.eval() with torch.no_grad(): print("预测值:", model(torch.tensor([[1.0]])).item())这里能明显感受到两种写法差异:TensorFlow 用model.fit封装好训练循环,PyTorch 需要自己写optimizer.zero_grad()、loss.backward()、optimizer.step()三步。对初学者来说,PyTorch 的显式写法反而更容易理解梯度是怎么更新的;对只想快速跑通模型的人来说,Keras 的封装更省事。
5.3 MNIST 手写数字识别对比
MNIST 是深度学习的经典 Hello World。这里不贴完整代码,只说明两个框架各自的关键差异:
- TensorFlow / Keras:用
model.fit一行启动训练,用model.evaluate在测试集上评估。 - PyTorch:需要手写 DataLoader 迭代训练,用
torch.nn.CrossEntropyLoss配合torch.optim.Adam。
从学习角度看,PyTorch 的写法更接近“底层原理”,能让你看到每个 batch 的 loss 变化;TensorFlow 的 Keras 封装更接近“开箱即用”。
判断一个训练是否成功,标准很简单:训练 loss 持续下降,验证集准确率稳定在 95% 以上(MNIST 任务)。如果 loss 不下降,优先检查学习率、数据归一化、网络层是否有误。
6. 接口 API 与批量任务生态
深度学习框架的“接口能力”和“批量任务”训练到一定规模之后就必须考虑。
6.1 TensorFlow 的工程化接口
TensorFlow 在部署侧的核心思路是“先导出 SavedModel,再交给 Serving 或 Lite”。
# 导出 SavedModel 示例 model.export('saved_model/1/')导出后,可以用 TensorFlow Serving 提供 gRPC / HTTP 接口。这类接口天然支持多版本模型管理和热加载,适合做线上推理服务。
批量训练方面,TensorFlow 提供了tf.data数据集管道,可以高效做 sharding、预取、并行解码。如果你要跑大规模分布式训练,tf.distribute.MirroredStrategy可以快速把单卡代码扩展到多卡。
6.2 PyTorch 的工程化接口
PyTorch 的部署有两条主流路线:
- 导出 ONNX,交给 ONNX Runtime 或 TensorRT 推理。
- 使用 TorchServe 提供 HTTP / gRPC 服务。
# 导出 ONNX 示例 torch.onnx.export(model, tensor_input, "model.onnx")相比 TensorFlow 的“全家桶”风格,PyTorch 的工程化更依赖社区生态组合。但好处是灵活,很多推理引擎都优先支持 PyTorch 导出的模型格式。
批量任务方面,PyTorch 的DataLoader配合num_workers参数可以并行加载数据。分布式训练则使用torch.distributed,在超大规模集群训练中与 TensorFlow 能力相当。
6.3 两者互转的注意事项
实际工作里经常遇到“PyTorch 训练、TensorFlow Serving 部署”的混合场景。一般做法是:
- PyTorch 导出 ONNX。
- 用 ONNX Runtime 直接推理。
- 必要时把 ONNX 转成 TensorFlow SavedModel 格式。
- 再用 TensorFlow Serving 部署。
这种跨框架链路能跑通,但调试成本和兼容性风险都不低。建议项目初期就明确训练和部署框架,尽量保持一致,避免后期迁移踩坑。
7. 资源占用与性能观察
深度学习框架的性能不能只看框架本身,还要看模型、数据、硬件和优化手段。这里给出一套通用的观察方法。
7.1 显存和内存怎么看
训练过程中,推荐用nvidia-smi实时观察显存占用。
watch -n 1 nvidia-smi在 Windows 上可以使用:
nvidia-smi -l 1观察时重点看两个字段:
- 显存占用:判断当前 batch size 和模型大小是否超过显存上限。
- 功耗和 GPU 利用率:利用率接近 90% 以上说明 GPU 在干活,如果利用率很低但显存占用很高,可能卡在数据加载。
代码里也可以打印当前显存占用:
import torch print(torch.cuda.memory_summary())7.2 CPU 推理和 GPU 推理
两个框架都支持 CPU 训练和推理,但日常数值差异很大:CPU 跑小模型可以接受,跑 CNN / Transformer 会非常慢。如果你只有 CPU,建议先减小数据集、降低 batch size、降低图像分辨率。
GPU 版本安装正确但代码里没有调用 GPU,也是常见情况。TensorFlow 中可以用tf.debugging.set_log_device_placement(True)打印设备分配,PyTorch 中要把模型和数据显式.to('cuda')。
7.3 如何降低显存占用
如果训练时遇到CUDA out of memory,按顺序尝试:
- 减小
batch_size,例如从 32 降到 16 或 8。 - 降低输入分辨率。
- 使用混合精度训练。PyTorch 用
torch.cuda.amp,TensorFlow 用tf.keras.mixed_precision.set_global_policy('mixed_float16')。 - 释放不再使用的中间变量。
- 使用梯度累积模拟更大 batch size。
显存占用没有固定值,不同模型差异巨大。比如一个简单全连接网络在 GTX 显卡上可能不到 1GB,而一个大型 Transformer 可能在 24GB 旗舰卡上都有压力。具体以你本机实测为准。
8. 常见问题与排查方法
| 问题现象 | 可能原因 | 排查方式 | 解决方案 |
|---|---|---|---|
| pip 安装下载慢或超时 | 网络问题 | 观察下载速度 | 使用国内 pip 镜像 |
| TensorFlow 检测不到 GPU | 驱动版本过低或 CUDA 运行时不匹配 | nvidia-smi查看驱动版本 | 更新显卡驱动;确认 TensorFlow 版本对应的 CUDA 要求 |
PyTorch 的.cuda()报错 | 安装的是 CPU 版 torch | torch.cuda.is_available()返回 False | 重新安装匹配 CUDA 的版本 |
CUDA out of memory | batch size 过大或模型过大 | nvidia-smi看显存占用 | 减小 batch size、降低分辨率、开混合精度 |
| 导入 PyTorch 报错找不到 DLL / .so | 缺少动态链接库或路径不对 | 检查 pip list 中 torch 版本 | 重装匹配系统环境的版本;确认 CUDA 驱动安装 |
| 模型权重加载报错 | PyTorch 2.6 的weights_only默认值改变 | 看报错是否为 UnpicklingError | 加载时显式设置weights_only=False |
| 训练 loss 不下降 | 学习率不合适、数据未归一化、模型结构错误 | 打印每个 epoch 的 loss | 调小学习率;对输入做标准化;简化模型先验证代码 |
| 端口被占用 | 启动了多个 Jupyter 或服务 | `netstat -ano | grep 8888` 查看端口 |
| conda 环境无法激活 | Shell 未初始化 | 重开终端 | 执行conda init bash |
这些是高频问题。遇到新报错时,先读完整报错信息,再复制关键词到搜索引擎,比凭感觉乱改参数要高效得多。
9. 最佳实践与选型建议
9.1 学生和入门用户
如果你是学生,正在上机器学习或深度学习课程,建议直接选 PyTorch。理由很实际:当前大量课程、论文、开源项目都在用 PyTorch,从 PyTorch 入门可以减少“复现代码时改框架”的成本。
学习顺序可以参考:
- 先掌握 PyTorch 基础张量操作。
- 用
nn.Module写一个多层感知机。 - 跑通 MNIST 或 CIFAR-10 分类。
- 阅读并复现一个 Transformer 分类模型。
- 学习 DataLoader 和训练循环的工程化写法。
9.2 工业部署和已有系统
如果你的目标是把模型部署到公司现有系统,TensorFlow 的价值在于配套工具齐全。从数据验证到模型导出、服务部署,都有官方方案。如果团队里已经有人维护 TensorFlow 服务,跟现有技术栈保持一致比折腾跨框架转换更稳妥。
9.3 两个框架都要学吗
不必一次性同时学。建议先精通一个,再通过 ONNX 等中间格式打通另一个。日常科研和开发中,你更多是“读别人的代码”,PyTorch 生态的代码量明显更多,所以从 PyTorch 起步的收益更高。
但如果你想做移动端 AI 或嵌入式设备上的推理,TensorFlow Lite 是绕不开的选项。掌握 PyTorch 的基本模型训练能力,再学习 TFLite 的转换和部署流程,是更实用的组合。
9.4 一些工程化习惯
- 用版本管理工具管理依赖,至少记录
pip freeze > requirements.txt。 - 模型权重、训练日志、测试数据分目录管理,避免全部堆在根目录。
- 跑批量训练任务时,写一个配置文件统一管理超参数,而不是每次改代码。
- 接口服务要限制外部访问,不要裸奔在公网。
- 涉及人脸、声音、版权数据时,必须确认数据来源合法,获得必要授权。
10. 总结与下一步
TensorFlow 和 PyTorch 都是成熟的深度学习框架,没有绝对的“最好”,只有是否符合你的场景。PyTorch 胜在研究、教育和开源模型生态,TensorFlow 胜在生产部署和端侧链路。但两者的边界正在不断模糊,未来跨框架转换也会越来越方便。
如果你刚开始,建议先按这套流程走一遍:
- 用 conda 创建虚拟环境。
- 安装 PyTorch,跑通线性回归。
- 跑通 MNIST 分类。
- 再安装 TensorFlow,跑同样的任务,感受两种写法的差异。
- 根据你的实际项目方向选择主力框架。
最容易踩的坑集中在三个地方:GPU 版本没装对、Python 版本不兼容、PyTorch 2.6 权重加载参数变化。把这几个问题提前规避掉,后面训练模型会顺利很多。
接下来可以继续扩展的方向包括:用 TensorBoard 可视化训练曲线、用torch.compile加速 PyTorch 模型、把模型导出为 ONNX 并接入 ONNX Runtime 推理、在 Jetson 等嵌入式设备上部署轻量化模型。无论选哪个框架,先跑通一个完整流程比纠结选型更重要。