PyTorch 和 TensorFlow 到底选哪个?这个问题在 CSDN 评论区每年都要吵一轮。我的建议是:别急着站队,两个都装上,用一个完整流程把环境搭建、GPU 验证、训练推理、模型导出全部跑通,用实际体验代替口舌之争。这篇文章就把 PyTorch 和 TensorFlow 的本地部署、功能验证、常见坑位一次性理清楚,覆盖从 Anaconda 虚拟环境到模型 API 服务的完整链路。
你能从这篇里拿到什么:一台机器上同时使用两套框架的隔离方案;CPU 版与 GPU 版安装差异;CUDA 是否真正被调用的验证方法;手写数字识别、图像分类、迁移学习、批量推理、ONNX 导出的最小可运行代码。文章后半段还整理了 RTX 5060 等 50 系显卡、Jetson 嵌入式平台、昇腾设备的安装注意事项,以及 PyTorch 2.6 之后torch.load默认参数变化带来的兼容性问题。看完可以直接照着操作。
1. 核心能力速览
先给一张对比表,把两个框架在本地部署和学习使用中的关键差异列清楚。以下信息来自公开社区资料和官方文档,具体版本特性请以实际环境为准。
| 对比项 | PyTorch | TensorFlow |
|---|---|---|
| 主要语言 | Python | Python |
| 图执行方式 | 动态图为主(Eager),可torch.compile静态优化 | 默认 Eager,支持tf.function静态图 |
| 安装方式 | pip / conda / 源码 | pip / conda / Docker |
| GPU 支持 | CUDA、ROCm、昇腾等 | CUDA、ROCm 等 |
| 典型应用 | 研究原型、Transformer、扩散模型、强化学习 | 工业部署、移动端、生产管道 |
| 部署方式 | TorchScript、ONNX、LibTorch | SavedModel、TensorRT、TF Lite、TF.js |
| 新手友好度 | 代码风格接近 NumPy,调试直观 | 接口多,历史版本差异大 |
| 社区趋势 | Hugging Face 等新研究模型基本首发 PyTorch | 生产案例和工具链积累深厚 |
从 2024 年社区讨论趋势看,学术开源项目向 PyTorch 倾斜明显,尤其在 Transformer、目标检测、强化学习这类方向,PyTorch 资料更新快、示例多。TensorFlow 的价值则更多体现在生产落地、端侧部署和与老系统集成。对于学习阶段,我的判断是两个框架都应该至少跑通一次,理解它们的编程范式和部署思路差异,比站队更有用。
2. 适用场景与使用边界
2.1 什么场景优先选 PyTorch
如果你是做研究实验、快速验证想法,或者主要参考 Hugging Face 上的模型仓库,PyTorch 更顺手。社区里热度很高的 PyTorch 手写数字识别、基于 PyTorch 的猫狗分类、PyTorch 目标检测、PyTorch 实现 Transformer 等教程,基本覆盖了入门到进阶的完整学习路线。PyTorch 动态图的调试体验接近普通 Python 程序,print中间张量、逐行断点都很自然,这对初学者非常友好。
2.2 什么场景优先选 TensorFlow
TensorFlow 的传统优势在工业部署。Keras 高层 API 封装度高,训练流程可以写得很简洁;SavedModel 格式和 TensorFlow Serving 配合稳定,端侧有 TF Lite、TF.js。如果公司存量系统使用 TensorFlow,或者项目需要在上线后做严格的服务化部署,TensorFlow 的工程案例和历史资料会帮到你。
2.3 不适合什么场景,以及合规边界
两个框架都不适合在无 GPU 环境下训练大模型,CPU 推理速度会明显偏慢。另外,不要在一个虚拟环境里同时安装两套框架的 GPU 版本,依赖冲突排查成本很高。
涉及图像、人脸、声音、版权素材的使用场景,必须先确认授权。比如用猫狗分类数据集可以,但不能随意爬取他人图片训练可能涉及肖像权的内容;用任何预训练模型都要检查其 License,尤其注意不能将非商用模型未经授权用于商业发布。本文所有示例代码仅用于本地环境学习和功能验证。
3. 环境准备与前置条件
以 Windows 10/11 或 Ubuntu 20.04/22.04 为例。核心原则是:先检查硬件和驱动,再决定安装版本,最后创建独立虚拟环境。
3.1 环境检查清单
# 1. 查看显卡驱动和可用 CUDA 版本 nvidia-smi # 2. 查看 Python 版本 python --version # 3. 查看 Conda 是否可用 conda --versionnvidia-smi输出右上角会显示驱动支持的 CUDA 版本号。这个数字代表驱动能支持的运行环境上限,不一定是自带 CUDA Toolkit。PyTorch 安装时更多看 wheel 包对应的 CUDA 版本,比如 cu121、cu124、cu128,需要与驱动兼容。
需要特别注意的是,RTX 5060 这类 Blackwell 架构 50 系显卡,通常需要较新的驱动和 CUDA 12.8 以上的运行环境支持,PyTorch 需要选择 cu128 或更新的 wheel 版本。老显卡如 GTX 1660、RTX 3060 则一般向下兼容。
3.2 Anaconda 环境隔离
PyTorch 和 TensorFlow 的依赖存在冲突风险,最稳妥的做法是建立两个独立环境:
conda create -n pytorch-env python=3.11 -y conda create -n tf-env python=3.11 -yPython 版本建议选择 3.9 到 3.12。从官方支持范围看,PyTorch 2.6 和 TensorFlow 2.18 都覆盖这一区间。选择 3.11 较为折中,兼容性和生态支持都比较好。
磁盘空间方面,每个虚拟环境加依赖大约需要 5 到 10 GB,模型缓存文件另算,建议预留 30 GB 以上。
3.3 国内镜像加速
安装依赖时建议使用清华 TUNA 或阿里云镜像源。Conda 配置:
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/ conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/ conda config --set show_channel_urls yespip 配置:
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple4. 安装部署与启动方式
4.1 PyTorch 安装
PyTorch 的 GPU 版本安装,最可靠的方式是打开 PyTorch 官网 Get Started 页面,选择当前系统生成安装命令,不要直接复制网上旧教程的命令。
以 CUDA 12.8 对应的 cu128 wheel 为例:
conda activate pytorch-env pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu128CPU 版本:
conda activate pytorch-env pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu如果使用国内 pip 镜像,需要注意镜像源里默认的 torch 可能是 CPU 版本,要确认安装包名称和来源。昇腾设备的用户需要到 Ascend 社区下载适配版本,官方 PyTorch wheel 无法直接使用;Jetson 平台则需要匹配 JetPack 版本,比如 JetPack 6.2.2 通常需要安装 NVIDIA 官方提供的预编译 PyTorch wheel,而不是 pip 直接安装。
4.2 TensorFlow 安装
TensorFlow 安装相对简单,CPU 和 GPU 的安装命令在 pip 层面已经统一:
conda activate tf-env pip install tensorflow这里有一个重大变化值得知道:TensorFlow 2.16 以后,Linux 上通过 pip 安装的tensorflow包会自动附带对应的 CUDA 和 cuDNN 依赖(以 nvidia-* 包形式),不再需要手动安装 CUDA Toolkit。这一点比 PyTorch 省事不少。Windows 原生环境对 GPU 的支持则依然不如 Linux 稳定,如果要在 Windows 上跑 TensorFlow GPU,WSL2 是更常见的选择。
4.3 验证安装是否成功
激活 PyTorch 环境,运行:
import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0))激活 TensorFlow 环境,运行:
import tensorflow as tf print(tf.__version__) print(tf.config.list_physical_devices('GPU'))PyTorch 中torch.cuda.is_available()返回True才代表 CUDA 可用;TensorFlow 中如果 GPU 设备列表不为空,说明已识别到显卡。
5. 功能测试与效果验证
安装只是第一步,真正判断环境是否正常,要用一个小型训练或推理任务做验证。下面给出几个可以在自己电脑上跑的最小示例。
5.1 PyTorch 手写数字识别(MNIST)
这是社区里出现频率最高的入门实战项目,适合验证 PyTorch 数据加载、模型定义、训练循环是否正常:
import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets, transforms transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_set = datasets.MNIST(root="./data", train=True, download=True, transform=transform) train_loader = DataLoader(train_set, batch_size=64, shuffle=True) class MLP(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(28 * 28, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = x.view(x.size(0), -1) x = torch.relu(self.fc1(x)) return self.fc2(x) model = MLP() optimizer = torch.optim.Adam(model.parameters(), lr=0.001) criterion = nn.CrossEntropyLoss() for epoch in range(2): for batch_idx, (data, target) in enumerate(train_loader): optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() if batch_idx % 200 == 0: print(f"epoch {epoch}, batch {batch_idx}, loss {loss.item():.4f}")预期输出是 loss 逐步下降。如果下载 MNIST 数据集较慢,可以提前下载后放到./data目录。
5.2 PyTorch 图像分类与迁移学习
猫狗分类是社区常见项目。做迁移学习时,需要理解冻结部分模型参数的写法:
import torchvision.models as models model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1) # 冻结 features 层,只训练最后一层 for name, param in model.named_parameters(): if "fc" not in name: param.requires_grad = False model.fc = torch.nn.Linear(model.fc.in_features, 2) optimizer = torch.optim.Adam(model.fc.parameters(), lr=0.001)冻结参数后,反向传播只会更新fc层参数,训练速度更快,显存占用也更低。这种策略在自定义数据集上非常常用。
5.3 PyTorch 目标检测验证
如果只想快速验证目标检测环境,可以直接用 torchvision 提供的预训练模型做推理,不需要先训练:
import torch from torchvision import transforms from PIL import Image model = torch.hub.load("ultralytics/yolov5", "yolov5s", pretrained=True) model.conf = 0.4 img = Image.open("test.jpg") results = model(img) results.show()YOLOv5 的 torch.hub 方式可以快速跑通检测流程。如果下载模型超时,可以手动下载权重文件放入缓存目录。
5.4 TensorFlow 图像分类验证
TensorFlow 侧用 Keras 做一个最简单的 CNN 分类,用 CIFAR-10 数据集验证整体链路:
import tensorflow as tf from tensorflow.keras import layers, models, datasets (x_train, y_train), (x_test, y_test) = datasets.cifar10.load_data() x_train, x_test = x_train / 255.0, x_test / 255.0 model = models.Sequential([ layers.Conv2D(32, (3, 3), activation="relu", input_shape=(32, 32, 3)), layers.MaxPooling2D((2, 2)), layers.Conv2D(64, (3, 3), activation="relu"), layers.MaxPooling2D((2, 2)), layers.Flatten(), layers.Dense(64, activation="relu"), layers.Dense(10) ]) model.compile(optimizer="adam", loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), metrics=["accuracy"]) model.fit(x_train, y_train, epochs=3, validation_data=(x_test, y_test))TensorFlow 的 Keras 高层 API 把训练循环进一步封装,代码量比 PyTorch 原生写法更少。但这也意味着调试时需要清楚底层在做什么,否则出了问题不容易定位。
5.5 PyTorch 实现 Transformer 小样例
如果需要验证 Transformer 结构,可以在自己的环境里实现一个最小化编码器块:
import torch import torch.nn as nn class MiniTransformerEncoderLayer(nn.Module): def __init__(self, d_model=64, nhead=4, dim_feedforward=128): super().__init__() self.self_attn = nn.MultiheadAttention(d_model, nhead, batch_first=True) self.linear1 = nn.Linear(d_model, dim_feedforward) self.linear2 = nn.Linear(dim_feedforward, d_model) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) def forward(self, x): attn_out, _ = self.self_attn(x, x, x) x = self.norm1(x + attn_out) ff_out = self.linear2(torch.relu(self.linear1(x))) x = self.norm2(x + ff_out) return x model = MiniTransformerEncoderLayer() x = torch.randn(2, 32, 64) print(model(x).shape)能从(2, 32, 64)变成(2, 32, 64),说明张量流动正常。做理解实验时,多打印shape能快速排查维度错误。
5.6 批量推理模板
实际项目中单张推理往往不够,需要批量处理一个目录下的所有图片。下面给出一个 PyTorch 批量推理模板,两个框架都可以借鉴这种设计:
import glob import json import torch from PIL import Image from torchvision import transforms model.eval() transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) image_dir = "./images" results = [] for img_path in sorted(glob.glob(f"{image_dir}/*.jpg")): try: img = Image.open(img_path).convert("RGB") tensor = transform(img).unsqueeze(0) with torch.no_grad(): output = model(tensor) pred = output.argmax(1).item() results.append({"path": img_path, "class_id": pred}) except Exception as e: results.append({"path": img_path, "error": str(e)}) with open("results.json", "w", encoding="utf-8") as f: json.dump(results, f, ensure_ascii=False, indent=2)批量任务的核心原则是:单张失败不能中断整体流程,错误信息要记录到日志文件,处理结果按路径或目录回写。
5.7 模型导出与格式转换
训练完成后,模型部署需要导出为通用格式。PyTorch 导出 ONNX:
dummy_input = torch.randn(1, 3, 224, 224) torch.onnx.export(model, dummy_input, "model.onnx", opset_version=17)TensorFlow 保存 SavedModel:
model.save("saved_model", save_format="tf")ONNX 是跨框架转换的中间格式,可以进一步转为 TensorRT 或在其他推理引擎中加载。
6. 接口 API 与批量任务落地
本地训练完成后,下一步通常是提供服务接口。这里给一个基于 FastAPI 的推理服务示例,逻辑对 PyTorch 和 TensorFlow 都适用。
from fastapi import FastAPI from pydantic import BaseModel import torch from torchvision import transforms from PIL import Image app = FastAPI() class PredictRequest(BaseModel): image_path: str transform = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), ]) model = torch.load("model.pth", map_location="cpu", weights_only=False) model.eval() @app.post("/predict") def predict(req: PredictRequest): img = Image.open(req.image_path).convert("RGB") tensor = transform(img).unsqueeze(0) with torch.no_grad(): pred = model(tensor) return {"class_id": pred.argmax(1).item()}启动服务:
uvicorn app:app --host 127.0.0.1 --port 8000用 curl 验证:
curl -X POST http://127.0.0.1:8000/predict \ -H "Content-Type: application/json" \ -d '{"image_path": "test.jpg"}'使用 Python requests 调用:
import requests url = "http://127.0.0.1:8000/predict" resp = requests.post(url, json={"image_path": "test.jpg"}, timeout=60) print(resp.json())接口服务上线时要注意:服务默认只监听本地地址更安全,如果部署到内网,需要增加鉴权;不要在公网暴露没有认证的推理服务。
批量任务方面,建议设计成「输入目录 + 输出目录 + 日志文件」的三段式结构。任务运行时把每张图片的处理状态写入日志,失败的任务单独记录原因,方便重跑。
7. 资源占用与性能观察
7.1 观察显存和 GPU 使用
训练或推理过程中,单独开一个终端观察 GPU 状态:
nvidia-smi -l 1这里重点看Memory-Usage和Volatile GPU-Util两列。如果显存涨到一定程度后不再变化,说明模型和数据加载已经达到稳定状态。
7.2 在代码中查看显存占用
PyTorch 中可以在训练循环里加入显存统计:
print(torch.cuda.memory_allocated() / 1024**3, "GB") print(torch.cuda.max_memory_allocated() / 1024**3, "GB")TensorFlow 默认会尽可能占用显存,可以通过按需增长让显存使用更灵活。这个设置在程序启动时必须放在第一次操作 GPU 之前执行:
import tensorflow as tf physical_devices = tf.config.list_physical_devices('GPU') if physical_devices: tf.config.set_memory_growth(physical_devices[0], True)7.3 性能影响因素
影响训练和推理速度的因素主要是:batch size、输入分辨率、模型参数量、迭代次数或训练 epoch。batch size 越大,单位时间处理的样本越多,但显存占用也越高;分辨率越大,前向计算量成倍增加;epoch 越多,训练总时间线性增长。
降低显存占用的常见做法包括:减小 batch size、使用混合精度训练(PyTorch 的torch.autocast、TensorFlow 的mixed_float16)、冻结部分层参数、关闭不需要的梯度计算(torch.no_grad或tf.stop_gradient)。
7.4 CPU 推理与 GPU 推理
CPU 可以跑推理,但速度会明显偏慢,尤其在 ResNet、Transformer 这类模型上。如果只是验证流程,在 CPU 上用小模型跑一遍没问题;如果做批量推理,还是建议在有 GPU 的环境执行。不同电脑的硬件差异很大,实际提速比例以本机测试为准。
8. 常见问题与排查方法
| 问题现象 | 可能原因 | 排查方式 | 解决方案 |
|---|---|---|---|
PyTorch 安装后torch.cuda.is_available()返回 False | 驱动版本过旧或 wheel 的 CUDA 版本与驱动不匹配 | nvidia-smi查看驱动,打印torch.version.cuda | 升级驱动,或安装与驱动匹配的 cu 版本 wheel |
| RTX 5060 等 50 系显卡无法识别 | 驱动版本太低或 PyTorch wheel 版本过旧 | 查看驱动版本,确认 PyTorch 是否支持 cu128 | 安装最新驱动,选择 cu128 或更新 wheel |
| 秋叶启动器 pytorch 安装总失败 | python 版本或依赖冲突、下载源不稳定 | 查看启动器日志,确认依赖安装到哪个包失败 | 手动创建虚拟环境,按官方命令安装后再用启动器 |
| TensorFlow 提示找不到 cudnn | TensorFlow 2.16 前需要手动安装 CUDA/cuDNN;2.16 后自动附带但环境变量可能异常 | 查看完整报错,检查 nvidia-* 依赖包 | Linux 下用 pip 重新安装 tensorflow;Windows 优先使用 WSL2 |
torch.load报错与weights_only相关 | PyTorch 2.6 起weights_only默认值改为 True,旧模型文件不符合安全加载要求 | 查看完整报错信息 | 使用torch.load(path, weights_only=False)加载,或重新保存模型 |
| 显存不足(CUDA out of memory) | batch size 过大、分辨率过高、模型过大 | nvidia-smi观察显存占用 | 减小 batch size,使用混合精度或梯度累积 |
| 两个框架装在同一环境互相干扰 | 依赖冲突,版本互相覆盖 | 用pip list查看冲突包 | 使用 conda 独立虚拟环境隔离 |
| Conda 创建环境缓慢 | 默认源网络慢 | 查看 conda 配置 | 配置清华 TUNA 镜像 |
| Jetson 平台 pip 安装 PyTorch 失败 | 普通 pip 源没有 Jetson 对应 wheel | 确认 JetPack 版本 | 从 NVIDIA 官方论坛或仓库下载匹配 PyTorch wheel |
| FastAPI 启动后端口被占用 | 端口已被其他进程使用 | `netstat -ano | findstr 8000` 查看占用 |
| 模型推理结果质量差 | 预训练数据分布与应用场景差异大 | 检查输入图像的预处理是否一致 | 使用迁移学习微调,或检查归一化参数 |
9. 最佳实践与使用建议
9.1 环境管理
- PyTorch 和 TensorFlow 必须使用独立虚拟环境,避免依赖冲突。
- 每个项目目录下保存
requirements.txt,记录关键包版本。 - 不要频繁升级大版本,切换版本前先记录当前可运行环境。
9.2 训练与调试
- 第一次运行用小 batch、小分辨率、少 epoch,先验证链路通不通。
- 训练结果不稳定时,优先固定随机种子,对比可复现性。
- 模型文件、数据集、输出结果分目录管理。
9.3 批量任务与接口服务
- 批量推理必须加日志和失败重试。
- 接口服务默认监听
127.0.0.1,部署到服务器时加鉴权和限流。 - 模型导出后要在目标环境做一次完整推理测试,避免转换后行为不一致。
9.4 数据与版权合规
- 训练数据来源要可追溯,预训练模型要检查 License。
- 涉及人脸、声音、版权素材,必须确认已获得合法授权。
- 不要将网上随意获取的模型和数据用于未经授权的商业场景。
10. 总结与下一步
这篇文章解决的是 PyTorch 和 TensorFlow 的「环境安装 + 功能验证 + 项目落地」问题。建议先装 PyTorch 跑通 MNIST,再装 TensorFlow 跑通 CIFAR-10 分类,最后做一次 ONNX 导出和 API 推理。两个框架都跑过一遍之后,你会发现它们解决的是同一类问题,只是接口习惯不同。
最容易踩的坑是版本匹配:RTX 5060 等新显卡需要新驱动和 cu128 的 PyTorch wheel;PyTorch 2.6 之后加载旧模型要关注weights_only参数;TensorFlow 在 Windows 原生环境建议走 WSL2。这些坑在动手安装之前就了解,能省下不少排查时间。
后续可以继续深入的方向:用 ONNX 做跨框架模型转换,在 PyTorch 里实现 Transformer 并对比 TensorFlow 的 Keras 写法,把训练好的模型接入 FastAPI 做成批量推理服务,或者基于两个框架分别完成一个目标检测项目的部署测试。装好环境只是起点,真正有价值的是你想解决的问题本身。