ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

TensorFlow与PyTorch全方位对比:安装、训练与选型指南

TensorFlow与PyTorch全方位对比:安装、训练与选型指南 TensorFlow 还是 PyTorch这个问题几乎每个刚接触深度学习的人都会遇到。这一期我们不看概念直接对比两大框架的安装路径、API 风格、模型训练体验、部署思路和资源占用给新手一个可以照着操作的选型参考。本文会演示 TensorFlow 和 PyTorch 在本地环境中的安装与验证流程并给出一个手写数字识别任务的训练示例帮你看清两者在实际开发中的差异。很多初学者最纠结的一点是不知道先学哪个。网上资料越看越混乱有人说 TensorFlow 工业界岗位多有人说 PyTorch 论文复现方便还有人说两个都得学。我的看法是先确定你的目标场景再选框架。如果毕业后想去做算法工程师、模型训练、论文复现PyTorch 是当前更顺的选择如果你所在团队已经有大量 TensorFlow 生产代码或者你需要接触移动端、服务端推理链路TensorFlow 仍然重要。接下来的内容会围绕两个框架展开先给一份核心能力速览再分别说明环境准备、安装方式、验证命令、一次完整训练、接口部署、资源占用和常见排错。新手可以按顺序操作一遍两边的真实手感就出来了。1. 核心能力速览下面先给两个框架做一个横向对比。注意具体数值和版本会随环境变化表格里写的是整体定位和当前主流状态实际以你本机测试为准。对比项TensorFlowPyTorch开源方Google 主导Meta 主导主要编程语言Python 为主另有 C、Java、JavaScript 等接口Python 为主底层 C当前主流版本2.x例如 2.18 属于较新版本线2.x例如 2.8 已支持较多新编译器能力API 风格高层用 Keras底层用 tf.GradientTape面向对象编程nn.Module autograd动态图/静态图2.x 默认 Eager 模式可用 tf.function 优化默认动态图torch.compile 可做编译优化模型部署TensorFlow Serving、TF Lite、TF.js 生态成熟TorchServe、ONNX Runtime、LibTorch新卡适配需要关注 CUDA 和 cuDNN 版本匹配新版驱动配合 CUDA 版本多数显卡可用社区活跃度老牌、企业存量多论文和开源项目的主流选择上手曲线概念略多早期静态图包袱重更接近写 Python容易理解从能力表可以看出两者已经不是“一个落后一个先进”的关系。TensorFlow 的强项在部署生态和跨端支持PyTorch 的强项在训练灵活度和社区资源。新手选型主要看你想走哪条路线。2. 两大框架的定位与生态差异2.1 TensorFlow 的发展路径TensorFlow 从诞生之初就一直强调“生产可用”。早期版本采用静态图模式先定义计算图再执行好处是性能优化空间大坏处是调试不直观新人经常被 graph 和 session 搞晕。到了 2.x官方把 Keras 设为默认高级 API默认 Eager 执行整体使用体验已经接近 PyTorch 那种“写普通 Python 代码”的感觉。TensorFlow 生态里值得关注的组件包括Keras 负责快速搭建网络TensorFlow Serving 负责模型上线TF Lite 负责移动端和嵌入式设备推理TF.js 负责浏览器端跑模型。如果你是奔着落地部署去的TensorFlow 的链路确实比较完整。2.2 PyTorch 的发展路径PyTorch 的设计思路非常直接用类封装网络层用 torch.Tensor 存数据用 autograd 自动求导。它的调试方式就是普通的 Python 调试方式你可以随时 print 中间张量可以在 forward 函数里打断点。这种设计让它在学术社区迅速普及现在很多论文的开源代码、预训练模型、HuggingFace 生态里的 Transformer 实现都默认用 PyTorch。PyTorch 2.x 引入 torch.compile 之后训练性能也有明显提升。配合 TorchServe 或 ONNX 导出PyTorch 模型同样可以走生产部署。从近年趋势看PyTorch 不再只是一个“研究框架”而是逐渐打通了训练到部署的路径。2.3 生态对比Keras、TorchVision、HuggingFace对于新手来说框架本身差别之外配套工具链会影响你的学习成本。TensorFlow 侧Keras 提供了一大堆现成的层、优化器、回调函数。配合 TensorBoard你可以很方便观察损失曲线和梯度分布。如果你做的是图像分类、回归、推荐系统这类常规任务用 Keras 几行代码就能搭一个 Baseline。PyTorch 侧torchvision 提供了常用数据集、预训练模型和图像变换工具。配合 HuggingFace Transformers你可以直接加载 BERT、GPT 等预训练模型做微调。现在如果你要复现 CV 或 NLP 论文大概率先搜到的就是 PyTorch 实现。3. 环境准备与前置条件不管安装 TensorFlow 还是 PyTorch前置环境大同小异。建议先准备好以下内容。3.1 Python 与 conda推荐安装 Anaconda 或 Miniconda用 conda 创建独立环境避免不同项目依赖冲突。Python 版本建议 3.9 到 3.11。TensorFlow 2.18 和 PyTorch 2.x 对这个范围的 Python 都有官方支持。安装命令如下实际以官方文档为准conda create -n deeplearning python3.10 -y conda activate deeplearning3.2 显卡驱动与 CUDA如果你使用 NVIDIA 显卡先确认驱动版本。Windows 下可以在命令行执行nvidia-smiLinux 下同样执行nvidia-smi查看右上角 CUDA Version。注意这个版本表示“当前驱动支持的最高 CUDA 版本”不等于你系统里已经安装了对应 CUDA Toolkit。安装 PyTorch 时只要装好显卡驱动然后通过 pip 安装带 cu118、cu121 或更新后缀的版本即可PyTorch 会自带所需 CUDA 运行库。TensorFlow 的 GPU 版也是类似思路pip 安装 tensorflow 时会拉取对应的 NVIDIA 依赖。如果你用的是新发布的 50 系显卡特别要注意驱动版本和 CUDA 版本的匹配。建议先更新到较新的显卡驱动再安装新版 PyTorch 或 TensorFlow。具体哪个组合能跑通以官方支持矩阵为准。3.3 磁盘空间与内存两个框架安装后的体积都不小。CPU 版本通常需要 1GB 到 3GB 磁盘GPU 版本加上 CUDA 相关依赖可能占 5GB 到 8GB。训练 MNIST 这类小模型8GB 内存够用训练大模型时要考虑数据加载时的内存占用。4. TensorFlow 安装部署与启动方式4.1 CPU 版安装如果你只是想先试框架没有独立显卡可以先装 CPU 版。CPU 版跑 MNIST 这样的任务也够用只是训练慢一些。conda activate deeplearning pip install tensorflow-cpu安装完成后打开 Python输出版本号import tensorflow as tf print(tf.__version__)能看到版本号说明安装成功。如果是 2.18 或相近版本说明基于较新的 2.x API。4.2 GPU 版安装GPU 版安装也很直接。先用nvidia-smi确认驱动正常然后直接安装完整版conda activate deeplearning pip install tensorflowTensorFlow 官方 PyPI 包会默认包含 GPU 支持的依赖。如果你的环境网络正常这一步会自动拉取 CUDA 相关库不需要手动去 NVIDIA 官网安装 CUDA Toolkit。4.3 验证 GPU 是否可用TensorFlow 安装完成后检查 GPU 是否被识别python -c import tensorflow as tf; print(tf.config.list_physical_devices(GPU))如果输出包含physical_device_desc和 GPU 名称说明 GPU 可用。如果输出为空需要检查驱动、CUDA 依赖和当前 TensorFlow 版本是否匹配。这一步是 TensorFlow 环境最容易卡住的地方。5. PyTorch 安装部署与启动方式5.1 CPU 版安装PyTorch 的 CPU 版安装命令非常简单conda activate deeplearning pip install torch torchvision torchaudio不过默认源会安装 CPU 版本吗不一定取决于你的 PyPI 源。最稳妥的方法是去 PyTorch 官网选择对应环境的命令。通常 CPU 版需要显式指定pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu5.2 GPU 版安装GPU 版需要指定 CUDA 版本后缀。例如你的驱动支持 CUDA 12.1可以这样安装conda activate deeplearning pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121新版驱动也可以尝试更新版本的 CUDA 后缀比如 cu124 或 cu126具体以 PyTorch 官网列出的 whl 版本为准。安装过程中如果下载速度不理想可以换国内镜像源但要注意镜像源里的 pypi 包是否完整。5.3 验证 PyTorch 和 CUDA安装完成后验证以下关键信息import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))如果第一行输出版本号第二行输出 True第三行输出你的显卡名称说明 PyTorch 已经能调用 GPU。如果cuda.is_available()是 False多半是 PyTorch 版本和 CUDA 库不匹配或者显卡驱动版本太旧。这里要多说一句网上经常有人反映“秋叶启动器 PyTorch 安装总失败”或者“5060 安装 PyTorch 失败”。通常原因是驱动版本落后或者安装的 whl 包与显卡架构匹配度不足。这类问题优先把驱动升到最新再选更新版本的 PyTorch 重新安装。6. 功能测试跑通一个手写数字识别任务基础环境准备好之后我用同一个任务——MNIST 手写数字分类——在两边各写一段代码。这个实验不追求精度重点看框架的代码组织方式和调试手感。6.1 TensorFlow 训练示例TensorFlow 2.x 推荐使用 Keras 快速搭建模型。完整代码如下import tensorflow as tf from tensorflow import keras # 加载 MNIST 数据集 (x_train, y_train), (x_test, y_test) keras.datasets.mnist.load_data() x_train, x_test x_train / 255.0, x_test / 255.0 # 搭建模型 model keras.Sequential([ keras.layers.Flatten(input_shape(28, 28)), keras.layers.Dense(128, activationrelu), keras.layers.Dense(10, activationsoftmax) ]) # 编译模型 model.compile(optimizeradam, losssparse_categorical_crossentropy, metrics[accuracy]) # 训练 model.fit(x_train, y_train, epochs3, batch_size32)代码很短适合新手快速看到结果。Keras 封装度很高你只需要定义层、编译、训练三步不需要手动写训练循环。这种“高抽象”对快速实验很友好但如果想自定义复杂训练逻辑反而需要绕过 Keras 的默认流程。6.2 PyTorch 训练示例PyTorch 需要显式定义网络类、损失函数、优化器和训练循环。代码如下import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms # 数据预处理 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) # 加载 MNIST train_dataset datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) # 定义网络 class Net(nn.Module): def __init__(self): super().__init__() self.fc nn.Sequential( nn.Flatten(), nn.Linear(28 * 28, 128), nn.ReLU(), nn.Linear(128, 10) ) def forward(self, x): return self.fc(x) model Net() optimizer optim.Adam(model.parameters()) loss_fn nn.CrossEntropyLoss() # 训练循环 for epoch in range(3): for images, labels in train_loader: optimizer.zero_grad() outputs model(images) loss loss_fn(outputs, labels) loss.backward() optimizer.step() print(fEpoch {epoch 1}, loss: {loss.item():.4f})这段代码比 TensorFlow 复杂一些但每个环节都是透明的数据加载、网络定义、前向传播、反向传播、参数更新全部能看能改。这也是 PyTorch 在研究者中更受欢迎的原因自定义损失函数、自定义训练逻辑、中间特征提取都很方便。6.3 两边训练感受对比从上面两个例子能明显看出风格差异。TensorFlow 代码更“快”几行就能跑。适合团队协作时快速迭代标准模型。但它把训练细节封装了新手看到model.compile和model.fit可能不理解内部发生了什么。PyTorch 代码更“显式”适合理解深度学习流程。新手从头写一次后对前向传播、反向传播、梯度清零这些概念会有更具体的感知。代价是代码量更多使用不熟悉的组件时需要自己查 API 文档。7. 接口 API 与批量任务框架选择不光影响训练还影响模型上线和批量推理。7.1 TensorFlow Serving 部署思路TensorFlow 模型可以用model.save(mnist_model, save_formattf)保存为 SavedModel 格式再用 TensorFlow Serving 启动服务# 启动一个模型仓库服务需要按实际模型路径调整 tensorflow_model_server --rest_api_port8501 --model_namemnist --model_base_path/models/mnist启动后可以通过 REST 接口发送图片数据做推理。这条链路在传统企业级部署中很成熟支持模型版本管理、热更新和多模型加载。7.2 TorchServe 部署思路PyTorch 官方推出了 TorchServe。把模型打包成.mar文件后用以下命令启动torchserve --start --model-store model_store --models mnistmnist.marTorchServe 提供 REST 接口也支持 gRPC并且有模型版本、批处理等能力。对于单个 PyTorch 模型的快速上线TorchServe 是比 Flask 自建服务更完整的选择。7.3 ONNX 通用中间方案如果你不想被框架绑死可以导成 ONNX。TensorFlow 有 tf2onnxPyTorch 自带torch.onnx.export。ONNX 模型可以用 ONNX Runtime 推理也可以接多个推理引擎。下面是 PyTorch 导出 ONNX 的示例import torch model Net() dummy_input torch.randn(1, 28, 28) torch.onnx.export(model, dummy_input, mnist.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}})导出后ONNX 模型可以部署到 Web 服务中。这种方式的好处是训练时用 PyTorch生产部署时用 ONNX Runtime两边解耦后续换框架影响也小。7.4 批量推理建议批量推理时需要注意预处理速度可能比 GPU 推理更慢。图片解码、缩放、归一化如果全在 CPU 上跑很容易形成瓶颈。建议把所有图片的预处理改成批量操作或者用 DataLoader、tf.data 这类管道来加速。另外批量任务要加日志和失败重试。推理服务偶发超时很常见建议给每个批次记录输入路径、开始时间、结束时间、结果状态失败任务单独落盘重新排队。8. 资源占用与性能观察8.1 显存观察方法训练时可以用 NVIDIA 显卡的工具观察显存占用。Windows 下用nvidia-smi或者任务管理器Linux 下直接执行watch -n 1 nvidia-smi注意显存占用不等于模型大小。模型权重、梯度、优化器状态都会占显存batch size 越大占得越多。MNIST 这种小模型显存占用很低通常几百 MB 到 1GB 左右而大模型可能出现 OOM需要调小 batch size 或改用混合精度。8.2 CPU 与 GPU 推理差异CPU 推理的优势是部署简单不需要独立显卡适合服务端低并发场景。GPU 推理在吞吐量和延迟上有优势尤其是同一个 Batch 里塞更多数据时GPU 优势更明显。新手的重点不是追求最快而是搞清楚当前任务的瓶颈数据加载慢、模型计算慢、还是输出后处理慢。一个实用的做法是分步打点测试。记录读图时间、预处理时间、推理时间、保存时间四个阶段哪个阶段耗时最长就优化哪里。很多时候瓶颈不在模型推理而在图片解码。8.3 降低显存占用的手段训练过程中如果遇到显存不足通常的调整顺序是先调小 batch size再降输入图像分辨率再考虑梯度累积。推理阶段可以关掉不需要的层比如 Dropout或者用半精度推理。实测时要注意不同框架对半精度支持的细节不同不一定所有算子都能跑 FP16。9. 常见问题与排查方法把网上新手遇到最多的问题整理成一张表方便排查。问题现象可能原因排查方式解决方案TensorFlow 安装后import tensorflow报错Python 版本过高或过低查看报错日志中的版本要求改用 3.9~3.11 创建环境PyTorch 安装后cuda.is_available()返回 FalsePyTorch whl 的 CUDA 版本和驱动不匹配执行nvidia-smi查驱动 CUDA 版本升级驱动或换对应 CUDA 后 缀的 PyTorch 包新显卡 50 系安装 PyTorch 失败驱动太旧或 PyTorch 版本不够新检查驱动版本和 PyTorch 版本装新驱动升级 PyTorch秋叶启动器类整合包难以定位问题依赖冲突、缓存残留查看完整报错堆栈清理缓存或重建虚拟环境显存不足 OOMbatch size 过大、分辨率过高用 nvidia-smi 观察显存占用减小 batch size、降分辨率、开混合精度训练时 GPU 利用率很低数据加载成为瓶颈查看 CPU 占用和 DataLoader 耗时增大 num_workers 或 prefetch端口被占用服务启动失败7860 / 8501 / 8080 被占用Linux 用ss -tlnpWindows 用netstat -ano换端口或杀进程模型导出 ONNX 时算子不支持框架版本、算子版本不一致查看导出报错中的算子名简化模型结构或升级框架版本批量推理结果不稳定图像预处理不一致、模型输入尺寸不一致检查缩放、归一化和 tensor 维度统一预处理流程最好做成同一个 pipeline10. 新手选型建议与总结先说结论如果你没有特殊的外部约束建议从 PyTorch 入门。因为它当前的资料最多、论文复现最方便、调试手感最好。入门阶段你最需要的是快速写代码、跑结果、理解训练流程PyTorch 在这一点的阻力最小。但这不代表 TensorFlow 不值得学。以下几种情况建议优先考虑 TensorFlow第一你所在公司或团队已有完整的 TensorFlow 技术栈后续维护和协作都要用 TF第二你明确要做嵌入式、移动端、浏览器端的模型部署TF Lite 和 TF.js 的生态更成熟第三你只是需要做标准模型的快速训练和上线Keras 的高层接口确实省事。最稳妥的路线是先用 PyTorch 把深度学习的核心概念弄明白比如张量、自动求导、优化器、训练循环再花一天时间过一遍 TensorFlow 的 Keras 接口。你会发现两者概念高度相似只是 API 命名和组织方式不同。之后再做项目时按需选择框架而不是被框架绑死。最后给新手一个实操顺序建议今天先装好一个框架把 MNIST 示例跑通明天跑第二个框架对比一下安装和训练的手感然后从一个你熟悉的任务开始比如猫狗分类、手写数字识别分别用两个框架各写一个版本。跑通之后再考虑模型保存、ONNX 导出、接口部署这些工程化问题。这一步走完你就不再是“该选哪个”的观望者而是两边都能动手的实践者了。
返回列表