PyTorch与TensorFlow双框架实战：从环境搭建到MNIST迁移

很多刚开始接触深度学习框架的读者都会在 PyTorch 和 TensorFlow 之间纠结很久。网上关于“到底选哪一个”的争论也特别多有人说 PyTorch 更适合研究代码写起来像原生 Python也有人说 TensorFlow 生产部署链路完善工业界用得稳。结果新人越看越焦虑甚至还没开始写模型就先被选型问题卡住了。这篇文章不打算帮你“二选一”。我反而想和你分享一条更务实的路线把 PyTorch 和 TensorFlow 都装起来都跑一遍然后从张量操作、自动求导、训练循环、模型保存这几个维度去对照理解。当你真正同时用过两个框架之后会发现它们之间的 API 差异只是表面底层的数据流思路、损失函数设计、梯度更新逻辑几乎是一致的。这个时候选型就不再是一个信仰问题而是一个根据业务场景做判断的工程问题。文章会从核心概念对比开始接着给出完整的 Anaconda 环境搭建方案再分别用 PyTorch 和 TensorFlow 实现一个手写数字识别案例最后整理高频报错的排查思路和工程化建议。内容偏实操所有代码我都会标注清楚你可以直接复制运行。1. 深度学习框架二选一成年人选择都要1.1 PyTorch 与 TensorFlow 各自解决了什么问题先看 PyTorch。PyTorch 由 Facebook AI Research 团队开发它的核心设计理念是“命令式优先”也就是动态计算图。什么叫动态计算图简单来说你在写模型的时候每一行张量运算都会立即执行遇到问题可以直接用 print 打印中间结果调试体验非常接近普通 Python 程序。这种特性让它在学术研究、快速原型验证、论文复现等场景里特别受欢迎。我们经常看到的 Transformer、Stable Diffusion、LLaMA 等开源模型的官方实现很多都基于 PyTorch。再看 TensorFlow。TensorFlow 由 Google 团队开发早期以静态计算图为主先定义整张计算图再送入会话执行。这种设计在分布式训练和线上部署方面有天然优势因为图结构一旦确定就可以做大量编译优化。不过静态图也让很多初学者觉得难上手所以 TensorFlow 2.x 之后默认开启了 Eager Execution也就是动态执行模式同时保留了 tf.function 和 SavedModel 这类面向生产环境的能力。再叠加 Keras 高层 APITensorFlow 现在也能让新手用十几行代码快速训练一个模型。所以要回答“谁更好”这个问题其实没有标准答案。PyTorch 在研究和灵活调试上有优势TensorFlow 在端到端生产链路和旧系统维护上有积累。两者都不是银弹但两者都值得你掌握基础用法。1.2 为什么建议两个框架都掌握第一开源社区和论文生态不会只围绕某一个框架转。你读论文时会发现有的作者开源 PyTorch 代码有的作者开源 TensorFlow 代码。如果你只会其中一个遇到另一个框架的优质项目就只能干瞪眼。第二实际工作中你可能会同时维护多个项目。很多公司早几年用 TensorFlow 搭建了推荐系统或 CV 服务现在新项目又切到 PyTorch作为开发人员你不能要求公司为了你的技术偏好把老项目全部重写。这时候“两个都会”就是刚需。第三框架之间的迁移能力本身就是一种核心竞争力。当你理解了张量、自动求导、优化器、数据管道这些概念之后从一个框架切到另一个框架只是换一套 API 的问题核心思想是可以平移的。下面我们就先来拆解这些核心概念。2. 三大核心概念对比张量、动态图与自动求导2.1 张量两个框架的共同基石不管 PyTorch 还是 TensorFlow最基础的数据结构都是张量。你可以把张量理解成“可以放到 GPU 上加速计算的 N 维数组”。0 维张量是标量1 维张量是向量2 维张量是矩阵3 维及以上通常代表批量数据或序列数据。PyTorch 中张量类型是torch.TensorTensorFlow 中张量类型是tf.Tensor。先看一段创建张量的对比代码# 文件路径tensor_demo.py import torch import tensorflow as tf # PyTorch 创建张量 a_torch torch.tensor([[1.0, 2.0], [3.0, 4.0]]) print(PyTorch 张量:\n, a_torch) print(设备:, a_torch.device) # TensorFlow 创建张量 a_tf tf.constant([[1.0, 2.0], [3.0, 4.0]]) print(\nTensorFlow 张量:\n, a_tf) print(设备:, a_tf.device)这里有个容易混淆的点PyTorch 用torch.tensor创建张量TensorFlow 用tf.constant创建常量张量。两者都支持shape、dtype、device这些属性。你还会看到 PyTorch 经常提到.to(device)TensorFlow 则通过tf.device或者自动分配策略管理设备。2.2 动态图与静态图调试体验差异的来源PyTorch 默认就是动态图模式。它每执行一行代码就计算一行所以你可以非常自然地在训练循环里打断点、打印中间张量。这也是为什么很多新手觉得 PyTorch 的代码“读起来像普通程序”。TensorFlow 2.x 默认也是 Eager 模式也就是动态执行。但 TensorFlow 完整保留了静态图能力你可以用tf.function装饰器把一段 Python 函数编译成静态计算图从而获得更好的执行性能。这种“既能动态调试又能静态加速”的设计是 TensorFlow 的特色。有一点需要提醒不要看到tf.function就盲目给所有函数加上。静态图编译对代码有约束比如不能随意使用 Python 的if和while依赖动态张量值否则很容易踩坑。刚开始学习时直接用默认的 Eager 模式就够了。2.3 自动求导训练模型的发动机深度学习的核心是反向传播而反向传播依赖自动求导。PyTorch 的自动求导基于autograd。默认情况下只有requires_gradTrue的张量会记录梯度。每次调用loss.backward()后梯度会累积到对应张量的.grad属性中。TensorFlow 的自动求导则通过tf.GradientTape()实现在with上下文中记录前向计算过程退出上下文后调用tape.gradient(loss, model.trainable_variables)拿到梯度。简单对比一下# 文件路径autograd_demo.py import torch import tensorflow as tf # PyTorch 自动求导 x torch.tensor(3.0, requires_gradTrue) y x ** 2 y.backward() print(PyTorch 梯度:, x.grad) # TensorFlow 自动求导 x_tf tf.Variable(3.0) with tf.GradientTape() as tape: y_tf x_tf ** 2 grad tape.gradient(y_tf, x_tf) print(TensorFlow 梯度:, grad.numpy())这里可以明显看出两套 API 的差异PyTorch 关注“张量是否记录梯度”TensorFlow 关注“用上下文管理器捕获哪些计算”。但本质都是链式法则的实现你只需要理解一次两边都能迁移。3. 环境准备从 Anaconda 到双框架安装3.1 安装 Anaconda 并创建虚拟环境推荐使用 Anaconda 管理 Python 环境因为它能很方便地创建独立虚拟环境避免不同项目之间的包冲突。打开终端输入下面的命令创建深度学习基础环境conda create -n dl python3.10 -y conda activate dl这里我用了 Python 3.10具体版本可以根据你的项目要求调整。如果项目依赖较老可能需要 Python 3.8 或 3.9如果跑最新的 Transformer 代码可以考虑 3.10 或 3.11。注意不同版本对依赖包的兼容性不同没有绝对最优按实际需求来。3.2 PyTorch 环境搭建与 GPU 安装思路安装 PyTorch 最准确的方式是打开 PyTorch 官网选择你的操作系统、包管理工具和 CUDA 版本官网会生成对应的安装命令。我把 CPU 版和 GPU 版都列一下方便你对比# CPU 版 pip install torch torchvision torchaudio # GPU 版示例命令具体 CUDA 版本以官网为准 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121安装 GPU 版之前强烈建议先用nvidia-smi查看显卡驱动支持的最高 CUDA 版本再根据 PyTorch 官方安装矩阵选择匹配的 CUDA 版本号。如果驱动版本较老强行装新版 CUDA 依赖可能会在运行时报“找不到 cudart64_*.dll”或“CUDA driver version is insufficient”之类的错误。如果你在国内网络环境下安装速度很慢可以使用清华 PyPI 镜像pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple安装完成后用下面这段代码验证 PyTorch 是否能正常调用 GPUimport torch print(torch.__version__) print(torch.cuda.is_available())如果torch.cuda.is_available()输出 True说明 GPU 环境可用。3.3 TensorFlow 环境搭建与版本差异TensorFlow 的安装命令相对简单CPU 版和 GPU 版在 2.x 中通常使用同一个包pip install tensorflowTensorFlow 在 Windows 上的 GPU 支持需要注意版本匹配。比如 TensorFlow 2.10 是最后一个原生支持 Windows GPU 的常规版本之后的 Windows GPU 支持方案发生了变化。如果你用的是较新的 TensorFlow 2.16 或更高版本建议先查看官方安装文档确认 CUDA 和 cuDNN 的版本要求。为了避免和 PyTorch 的依赖冲突我强烈建议你为 TensorFlow 单独创建一个虚拟环境conda create -n tf python3.10 -y conda activate tf pip install tensorflow创建一个独立环境并不是小题大做。PyTorch 和 TensorFlow 都依赖一些底层库比如numpy、protobuf、absl-py两者对版本的要求不完全一致。放在同一个环境里装完一个再装另一个很容易出现“装好了 AB 却启动报错”的情况。用虚拟环境隔离是最省心的方案。3.4 两个框架装在同一台机器上的协调方案如果你电脑上只有一个 GPU也不需要同时在一段代码里导入两个框架那么分开建两个虚拟环境是最稳妥的。日常使用时先conda activate dl就进入 PyTorch 环境再conda activate tf就切到 TensorFlow 环境。两个环境的 pip 包互不干扰这是我在多台机器上验证过最推荐的方案。如果你确实需要在同一个 Python 进程中同时使用两个框架也不是完全不可能但你需要非常小心依赖版本冲突。实际工程里这种需求并不多见通常只有做框架对比实验或者写迁移工具时才会遇到。因此不要为了省一点磁盘空间强行合并环境。4. 实战案例用 PyTorch 完成手写数字识别环境准备好之后我们来跑第一个完整案例。手写数字识别是深度学习的“Hello World”数据集是 MNIST包含 60000 张训练图片和 10000 张测试图片每张图片是 28×28 的灰度图。4.1 数据准备与 DataLoaderPyTorch 中我们通常使用torchvision下载 MNIST并用DataLoader把数据打包成 batch。完整代码如下# 文件路径mnist_pytorch.py import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms # 1. 数据预处理转张量 归一化 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) # 2. 下载并加载 MNIST 数据集 train_dataset datasets.MNIST( root./data, trainTrue, transformtransform, downloadTrue ) test_dataset datasets.MNIST( root./data, trainFalse, transformtransform, downloadTrue ) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size64, shuffleFalse) print(f训练集大小: {len(train_dataset)}) print(f测试集大小: {len(test_dataset)}) print(f每个 batch 大小: {64})这里有两个关键点。第一ToTensor()会把 PIL 图片转换成 0 到 1 之间的张量然后再用Normalize减去均值 0.1307、除以标准差 0.3081这是 MNIST 数据集的常用标准化参数。第二DataLoader会自动把数据切分成 batch训练时shuffleTrue能打乱样本顺序避免模型学习到样本顺序带来的偏差。4.2 定义神经网络与训练流程MNIST 图片是 28×28 的灰度图我们可以先用一个简单的全连接网络来理解训练流程。这里我把图片拉平成 784 维向量经过两个隐藏层最后输出 10 个类别的分数。完整代码如下# 文件路径mnist_pytorch.py继续 # 3. 定义简单的全连接神经网络 class SimpleNN(nn.Module): def __init__(self): super(SimpleNN, self).__init__() self.fc1 nn.Linear(28 * 28, 256) self.fc2 nn.Linear(256, 128) self.fc3 nn.Linear(128, 10) self.relu nn.ReLU() def forward(self, x): x x.view(x.size(0), -1) # 展平图片 x self.relu(self.fc1(x)) x self.relu(self.fc2(x)) x self.fc3(x) return x # 4. 初始化模型、损失函数、优化器 device torch.device(cuda if torch.cuda.is_available() else cpu) model SimpleNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 5. 训练循环 epochs 5 for epoch in range(epochs): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to(device) optimizer.zero_grad() outputs model(images) loss criterion(outputs, labels) loss.backward() optimizer.step() running_loss loss.item() avg_loss running_loss / len(train_loader) print(fEpoch [{epoch 1}/{epochs}], Loss: {avg_loss:.4f}) # 6. 测试集评估 model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(device), labels.to(device) outputs model(images) _, predicted torch.max(outputs, 1) total labels.size(0) correct (predicted labels).sum().item() print(f测试集准确率: {100 * correct / total:.2f}%)这段代码里有几个值得强调的习惯训练前调用model.train()测试前调用model.eval()。这会切换 BatchNorm 和 Dropout 等层的行为。optimizer.zero_grad()必须放在每次反向传播之前否则梯度会累加。测试时用torch.no_grad()包裹表示不需要计算梯度能减少内存占用和加速推理。4.3 运行结果说明正常情况下训练 5 个 epoch 后测试集准确率可以达到 97% 左右。如果使用 CNN 和更多训练轮次准确率可以轻松超过 99%。输出大概是这样Epoch [1/5], Loss: 0.3241 Epoch [2/5], Loss: 0.1453 Epoch [3/5], Loss: 0.1037 Epoch [4/5], Loss: 0.0785 Epoch [5/5], Loss: 0.0622 测试集准确率: 97.32%如果你在自己机器上跑出来的准确率略低或略高都很正常因为初始化权重具有随机性。只要损失在逐步下降说明模型正在正常学习。5. 实战案例用 TensorFlow/Keras 完成手写数字识别接下来我们用 TensorFlow 再实现一次完全相同的任务。你会发现 TensorFlow 的 Keras 高层 API 写起来更加“声明式”很多流程都被封装好了。5.1 加载与预处理数据TensorFlow 内置了 MNIST 数据集不需要额外下载脚本# 文件路径mnist_tensorflow.py import tensorflow as tf from tensorflow.keras import layers, models # 1. 加载 MNIST 数据 (x_train, y_train), (x_test, y_test) tf.keras.datasets.mnist.load_data() # 2. 归一化到 [0, 1] x_train x_train.astype(float32) / 255.0 x_test x_test.astype(float32) / 255.0 # 3. 展平 28x28 图片为 784 维向量 x_train x_train.reshape(-1, 28 * 28) x_test x_test.reshape(-1, 28 * 28) print(f训练集形状: {x_train.shape}) print(f测试集形状: {x_test.shape})注意 TensorFlow 的tf.keras.datasets.mnist.load_data()返回的是 NumPy 数组不需要显式转换成张量模型训练时会自动转换。5.2 构建模型并训练使用 Keras Sequential 模型结构与 PyTorch 示例保持一致输入 784 维两个隐藏层分别是 256 和 128输出 10 个类别。# 文件路径mnist_tensorflow.py继续 # 4. 构建模型 model models.Sequential([ layers.Dense(256, activationrelu, input_shape(784,)), layers.Dense(128, activationrelu), layers.Dense(10, activationsoftmax) ]) # 5. 编译模型 model.compile( optimizeradam, losssparse_categorical_crossentropy, metrics[accuracy] ) # 6. 训练模型 history model.fit( x_train, y_train, batch_size64, epochs5, validation_data(x_test, y_test) ) # 7. 测试集评估 test_loss, test_acc model.evaluate(x_test, y_test, verbose0) print(f测试集准确率: {test_acc * 100:.2f}%)这里有一个容易让新手混淆的地方losssparse_categorical_crossentropy和losscategorical_crossentropy有什么区别简单来说如果你的标签是整数形式比如 0、1、2、3……就使用sparse_categorical_crossentropy如果你的标签是 One-Hot 编码比如[0, 0, 1, 0]就使用categorical_crossentropy。MNIST 默认标签是整数所以要用sparse版本。5.3 模型结构与训练输出model.fit()会在每个 epoch 结束后打印训练损失、训练准确率、验证损失和验证准确率Epoch 1/5 938/938 [] - 3s 3ms/step - loss: 0.2205 - accuracy: 0.9349 - val_loss: 0.1070 - val_accuracy: 0.9672 Epoch 2/5 938/938 [] - 3s 3ms/step - loss: 0.0867 - accuracy: 0.9733 - val_loss: 0.0768 - val_accuracy: 0.9761 ... 测试集准确率: 97.50%938/938表示每个 epoch 有 938 个 batch。因为训练集有 60000 张图片除以 batch_size 64 得到的 batch 数量就是 938 左右。6. 双框架迁移思路与高频坑位排查6.1 从 PyTorch 到 TensorFlow 如何迁移当你明白一个训练流程包含“数据加载、模型定义、损失函数、优化器、训练循环”五个部分后两个框架的映射关系就很清晰了。我把对应关系整理成一张表功能模块PyTorchTensorFlow / Keras数据加载DataLoaderDatasettf.data.Dataset或 NumPy 数组模型定义继承nn.Module实现forwardSequential或继承keras.Model损失函数nn.CrossEntropyLoss等losses.SparseCategoricalCrossentropy等优化器optim.Adamoptimizers.Adam训练循环手动写for循环model.fit()或自定义train_step梯度计算loss.backward()tf.GradientTapePyTorch 的优势是流程透明每一步都手动控制适合学习和调试TensorFlow/Keras 的优势是封装度高你用几行代码就能完成同样的训练适合快速搭建和工程部署。两者没有绝对的优劣理解底层机制之后切换成本会很低。6.2 常见问题排查表问题现象常见原因解决思路torch.cuda.is_available()返回 FalsePyTorch 安装了 CPU 版或 CUDA 版本不匹配重新根据官网选择对应 CUDA 版本安装 GPU 版TensorFlow 找不到 GPUCUDA/cuDNN 版本与 TensorFlow 不匹配查看官方“Software requirements”文档确认驱动和 CUDA 版本运行时报CUDA driver version is insufficient显卡驱动过旧升级显卡驱动或回退到低版本 CUDA 的框架包pip 安装超时网络不稳定使用国内镜像源比如清华、阿里云镜像同时安装两个框架后 import 报错protobuf、numpy 等底层依赖冲突拆分成独立 conda 虚拟环境不要混装PyTorch 加载模型报WeightsOnlyUnpickler错误PyTorch 2.6 后torch.load默认weights_onlyTrue保存模型时改用state_dict加载时用load_state_dict或显式设置weights_onlyFalse这里我想单独展开说一下 PyTorch 2.6 的weights_only变化。从前很多教程习惯用torch.save(model, model.pt)保存整个模型对象再用torch.load(model.pt)直接加载。但在 PyTorch 2.6 中官方把torch.load的weights_only参数默认值改成了True这意味着加载时必须更多限制原本包含自定义类对象的 checkpoint 可能报错。更稳妥的做法永远是只保存模型参数# 保存 torch.save(model.state_dict(), model_weights.pt) # 加载 model SimpleNN() model.load_state_dict(torch.load(model_weights.pt)) model.eval()如果你一定要兼容旧代码需要显式设置weights_onlyFalse但要了解这会引入潜在的安全风险。因此在处理训练好的模型时尽量遵循官方推荐的新方式。6.3 项目中的选型判断清单如果你负责一个新项目可以从下面几个维度判断研发团队对哪个框架更熟悉。框架只是工具团队能快速交付更重要。项目是否需要快速原型验证、频繁改模型结构。如果是PyTorch 的调试体验更好。项目是否要部署到 TF Serving、Android 或旧版生产系统。如果是TensorFlow 的生态更成熟。项目是否需要复现最新论文方法。大多数最新论文开源代码是 PyTorch。项目是否已经有大量历史代码资产。不要为了追新而忽略维护成本。在团队没有明确技术栈约束时我一般建议科研实验、算法验证优先考虑 PyTorch端到端生产服务、已有 TensorFlow 体系优先考虑 TensorFlow。两边都装了两边都会写这才是“我全都要”的正确姿势。7. 学习路线与工程化最佳实践7.1 新手进阶路线如果你刚开始学不要一上来就追各种花哨的模型。建议按下面这条主线走掌握张量操作包括创建、切片、广播、矩阵乘法。掌握自动求导和梯度更新理解反向传播的代码表达。用全连接网络跑通 MNIST 或 CIFAR-10。学习卷积神经网络理解卷积核、池化、感受野。学习循环神经网络或 Transformer理解序列建模。尝试目标检测、图像分割、文本分类等实际任务。每到一个阶段最好都用 PyTorch 和 TensorFlow 各实现一遍。同一份模型结构两个框架各写一次你对框架的理解会立刻上升到“迁移”层面。7.2 工程化建议开发调试阶段可以随意探索但一旦要进入项目交付有几个习惯值得养成锁版本。用pip freeze requirements.txt锁住当前环境依赖避免未来升级导致结果不可复现。数据管道独立。不要把数据下载、预处理和模型训练耦合在同一个脚本里尽量拆成可复用的模块。定期保存 checkpoint。训练时间较长时建议每个 epoch 保存一次权重方便中断后从断点恢复。日志要结构化。除了打印 loss还要记录学习率、当前 epoch、GPU 显存占用、数据加载耗时等信息方便定位瓶颈。不要在训练代码里混入测试代码。用model.train()和model.eval()明确区分模式避免 BatchNorm 和 Dropout 行为异常。GPU 资源要做好监控。训练大模型时显存溢出是常态可以先降低 batch_size 验证代码逻辑再逐步增大。7.3 最后的建议框架之争在深度学习领域会一直存在但真正能提升你竞争力的不是“我选对了框架”而是“我能根据问题选择合适的工具”。如果你刚开始入门就按照这篇文章的路线把 PyTorch 和 TensorFlow 都装上都跑一遍 MNIST。当你发现两套代码只是在 API 层面不同底层逻辑完全相通时你就已经跨过了新手期最大的门槛。如果这篇文章对你有帮助建议先收藏备用。遇到安装报错或者训练异常时可以对照本文第 6 章的表格快速排查。后续我还会更新更多关于模型部署、Transformer 实战和 GPU 调优的内容欢迎持续关注。