之前在业务里同时维护两个算法项目时我经常被“PyTorch 还是 TensorFlow”这个问题卡住。网上观点各执一词有人说 PyTorch 学术生态无敌有人说 TensorFlow 部署链路成熟。实际上一个算法工程师如果只会其中一套遇到工业落地或者复现最新论文时会非常吃亏。本文不打算替你站队而是把两套框架放到同一套学习路径里完整串一遍从环境准备、核心 API 对比到 MNIST 手写数字识别的完整实战再给出安装失败、模型加载报错等高频问题的排查思路。无论你是刚入门深度学习还是已经能跑通简单模型但想补齐另一套框架这篇文章都能帮你省下不少查资料的时间。1. 背景与核心概念1.1 深度学习框架是什么为什么要同时掌握两个深度学习框架的本质是对“张量运算 自动求导 神经网络模块”的统一封装。框架存在的意义是让你把精力花在模型结构、数据处理、训练策略上而不是每次训练都从零写一遍反向传播。很多初学者容易陷入“二选一”的纠结但真实工程场景往往不是单选题。你要复现一篇最新论文开源实现大概率在 PyTorch 生态里你要给已有业务线做模型推理服务线上系统可能早就跑在 TensorFlow Serving 上。这种情况下能够同时看懂、同时操作两套框架会比单纯站队某一方更游刃有余。1.2 PyTorch 与 TensorFlow 的设计哲学差异PyTorch 由 Meta原 Facebook主导开发核心设计是动态计算图也就是“边运行边建图”。你写的前向代码会真实执行PyTorch 在背后通过自动微分记录梯度。这种模式的好处是调试非常直观你可以在任意一行打印中间张量可以使用原生 Python 的if、for、print甚至可以打一个断点进入pdb调试。TensorFlow 由 Google 开发经历了 1.x 静态图到 2.x 动态图的大转变。TensorFlow 2 默认也开启了 Eager Execution但保留了tf.function这套“编译为静态图”的能力适合对性能有高要求的训练和推理场景。同时TensorFlow 的高层 API 也就是 KerasSequential和Functional接口能让开发者快速搭出常见网络。1.3 两者并不是非此即彼从能力边界来说PyTorch 同样可以做部署TensorFlow 同样适合研究和教学。很多大型项目实际上是混合使用模型在 PyTorch 里训练和验证训练完成后导出为 ONNX 格式再接入 TensorRT 或 TensorFlow 推理链路。理解两套框架的 API 设计反而能让你更清楚地判断某个功能到底应该在哪套生态里完成。2. 环境准备与版本说明2.1 安装前的整体规划环境冲突是深度学习新手最常见的问题。PyTorch 和 TensorFlow 依赖的 CUDA 组件、底层库并不完全一致如果直接装进同一个 Python 环境很容易出现“装好这个、另一个就无法 import”的情况。最稳妥的做法是使用 Anaconda 分别创建两个独立环境pytorch_env - Python 3.10 PyTorch 2.x torchvision tf_env - Python 3.10 TensorFlow 2.x隔离之后两个环境互不干扰即使一个环境被装坏也不影响另一个。本文示例基于 Windows 10/11 或 Ubuntu 20.04/22.04Python 3.10。你的机器 CUDA 版本可能不同安装命令会有差异但排查思路是通用的。2.2 PyTorch 环境搭建先创建环境conda create -n pytorch_env python3.10 -y conda activate pytorch_env接着检查 GPU 驱动信息nvidia-smi上半部分输出里的CUDA Version: 12.x指的是驱动能够支持的最高 CUDA 版本安装 PyTorch 时可以选择等于或低于它的 CUDA 运行时版本。例如驱动支持 CUDA 12.1你可以安装 cu121也可以选更保守的 cu118。PyTorch 官网会根据你的系统环境生成安装命令。以 cu121 为例pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121如果官方源下载较慢可以换用国内镜像源例如清华源、阿里源。CPU 版本安装更简单pip install torch torchvision torchaudio这里有一个值得注意的坑创建环境时 Python 版本不要太新。某些 PyTorch 版本尚未适配最新的 Python 小版本会导致 pip 找不到对应 wheel 包。目前 Python 3.9 到 3.11 在大多数 PyTorch 版本中都比较稳妥。2.3 TensorFlow 环境搭建TensorFlow 安装比 PyTorch 特殊一点尤其是 Windows 平台。先创建独立环境conda create -n tf_env python3.10 -y conda activate tf_envCPU 版本pip install tensorflow-cpu如果是在 Linux 上想获得 GPU 支持可以安装pip install tensorflow需要注意从 TensorFlow 2.11 开始Windows 原生 pip 不再直接提供 GPU 支持。想在 Windows 上使用 GPU 运行较新版本的 TensorFlow比如 2.16、2.18 等推荐通过 WSL2Windows Subsystem for Linux或 Docker 安装。如果你的项目必须用 Windows 原生环境那么 TensorFlow 2.10 是最后支持原生 GPU 的版本。很多人装完 TensorFlow 后一运行就报 DLL 或 GPU 相关错误基本都是这个原因。Jetson 这类 ARM 嵌入式平台要特别注意不要直接从 PyPI 安装torch大概率没有对应架构的 wheel。Jetson 上通常使用 NVIDIA 针对 JetPack 系统编译的预编译包安装前先确认 JetPack 版本与 CUDA 版本再到 NVIDIA 官方资源中找匹配的.whl。2.4 安装后的检查清单安装完成后分别进入两个环境检查。PyTorch 检查import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU only)TensorFlow 检查import tensorflow as tf print(tf.__version__) print(tf.config.list_physical_devices(GPU))如果torch.cuda.is_available()返回False优先检查驱动版本、PyTorch 对应的 CUDA 运行时版本是否匹配以及当前虚拟环境里是否真的安装了 GPU 版本。3. 核心 API 与编程范式对比3.1 张量操作两个框架都以张量为基本数据结构只是类型名不同。PyTorch 是torch.TensorTensorFlow 是tf.Tensor底层都支持 GPU 运算和自动微分。下面分别创建两个张量# PyTorch import torch a torch.tensor([1.0, 2.0, 3.0]) b a * 2 print(b)# TensorFlow import tensorflow as tf a tf.constant([1.0, 2.0, 3.0]) b a * 2 print(b)在 PyTorch 中默认设备是 CPU要用 GPU 必须手动.to(cuda)TensorFlow 检测到 GPU 后会自动尝试使用 GPU。两者的哲学差异在这里也能看出来PyTorch 把设备管理交给开发者更显式TensorFlow 更自动但也容易让人忽略资源分配细节。3.2 模型定义PyTorch 推荐继承torch.nn.Module并在__init__中定义子层在forward中定义前向传播import torch.nn as nn class MyNet(nn.Module): def __init__(self): super().__init__() self.fc nn.Linear(128, 10) def forward(self, x): return self.fc(x)TensorFlow 中既可以写函数式模型也可以继承tf.keras.Modelimport tensorflow as tf from tensorflow.keras import layers, models # 方式一Sequential model models.Sequential([ layers.Dense(128, activationrelu), layers.Dense(10, activationsoftmax) ]) # 方式二子类化 class MyNet(tf.keras.Model): def __init__(self): super().__init__() self.fc1 layers.Dense(128, activationrelu) self.fc2 layers.Dense(10, activationsoftmax) def call(self, x): x self.fc1(x) return self.fc2(x)子类化时TensorFlow 的call方法对应了 PyTorch 的forward。调用方式都是model(x)只是内部逻辑分别走forward或call。3.3 训练循环PyTorch 的训练循环通常手动编写非常透明optimizer.zero_grad() outputs model(inputs) loss loss_fn(outputs, labels) loss.backward() optimizer.step()TensorFlow 2 使用tf.GradientTape()显式记录求导过程with tf.GradientTape() as tape: predictions model(inputs) loss loss_fn(labels, predictions) gradients tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables))两者本质相同前向传播、计算损失、反向传播、更新权重。区别在于 PyTorch 把求梯度隐藏在backward()内部TensorFlow 则通过tape.gradient()把“求梯度”这一步显式交给你。3.4 数据集与数据管道PyTorch 的数据处理通常基于torch.utils.data.Dataset和DataLoader你可以自定义 Dataset 类用迭代器批量取数。TensorFlow 推荐使用tf.data.Dataset它能把文件读取、预处理、乱序、批处理、预取组合成一条完整的数据管道。在中小型项目里两套体系差别不大。但在大数据量或生产链路中tf.data提供了更丰富的并行预处理能力PyTorch 则更强调和 Python 生态的无缝衔接方便做复杂的在线数据增强。3.5 PyTorch 2.6 权重加载变化从 PyTorch 2.6 开始torch.load的weights_only参数默认值变成了True。这个改动是为了提升安全性避免恶意 pickle 文件在反序列化时执行任意代码。带来的直接影响是一些旧代码直接torch.load(model.pth)会报WeightsUnpickler相关异常。如果你加载的是完全可信的模型权重可以显式设置weights_onlyFalsestate_dict torch.load(model.pth, map_locationcpu, weights_onlyFalse)如果加载的是标准state_dict默认的weights_onlyTrue通常能够直接工作。遇到这类报错时先判断模型文件来源再决定是否关闭限制。4. 完整实战——手写数字识别两种实现这里用 MNIST 手写数字识别做例子。原因是它对算力要求很低两套框架都内置了数据下载接口加上代码量不大正好可以同题对比。4.1 项目结构dl-stack/ ├── mnist_pytorch.py ├── mnist_tensorflow.py └── README.md两个脚本互相独立分别在自己的 conda 环境中运行。4.2 PyTorch 版 CNN 完整代码文件路径mnist_pytorch.pyimport 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,)) ]) train_dataset datasets.MNIST(root./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(root./data, trainFalse, downloadTrue, transformtransform) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size256, shuffleFalse) # 2. 定义 CNN 模型 class CNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 32, kernel_size3, padding1) self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) self.pool nn.MaxPool2d(2, 2) self.relu nn.ReLU() self.fc1 nn.Linear(64 * 7 * 7, 128) self.fc2 nn.Linear(128, 10) def forward(self, x): x self.pool(self.relu(self.conv1(x))) # 28 - 14 x self.pool(self.relu(self.conv2(x))) # 14 - 7 x x.view(-1, 64 * 7 * 7) x self.relu(self.fc1(x)) return self.fc2(x) model CNN() device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 3. 训练 for epoch in range(5): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to
PyTorch与TensorFlow双框架实战：环境配置到MNIST识别
之前在业务里同时维护两个算法项目时我经常被“PyTorch 还是 TensorFlow”这个问题卡住。网上观点各执一词有人说 PyTorch 学术生态无敌有人说 TensorFlow 部署链路成熟。实际上一个算法工程师如果只会其中一套遇到工业落地或者复现最新论文时会非常吃亏。本文不打算替你站队而是把两套框架放到同一套学习路径里完整串一遍从环境准备、核心 API 对比到 MNIST 手写数字识别的完整实战再给出安装失败、模型加载报错等高频问题的排查思路。无论你是刚入门深度学习还是已经能跑通简单模型但想补齐另一套框架这篇文章都能帮你省下不少查资料的时间。1. 背景与核心概念1.1 深度学习框架是什么为什么要同时掌握两个深度学习框架的本质是对“张量运算 自动求导 神经网络模块”的统一封装。框架存在的意义是让你把精力花在模型结构、数据处理、训练策略上而不是每次训练都从零写一遍反向传播。很多初学者容易陷入“二选一”的纠结但真实工程场景往往不是单选题。你要复现一篇最新论文开源实现大概率在 PyTorch 生态里你要给已有业务线做模型推理服务线上系统可能早就跑在 TensorFlow Serving 上。这种情况下能够同时看懂、同时操作两套框架会比单纯站队某一方更游刃有余。1.2 PyTorch 与 TensorFlow 的设计哲学差异PyTorch 由 Meta原 Facebook主导开发核心设计是动态计算图也就是“边运行边建图”。你写的前向代码会真实执行PyTorch 在背后通过自动微分记录梯度。这种模式的好处是调试非常直观你可以在任意一行打印中间张量可以使用原生 Python 的if、for、print甚至可以打一个断点进入pdb调试。TensorFlow 由 Google 开发经历了 1.x 静态图到 2.x 动态图的大转变。TensorFlow 2 默认也开启了 Eager Execution但保留了tf.function这套“编译为静态图”的能力适合对性能有高要求的训练和推理场景。同时TensorFlow 的高层 API 也就是 KerasSequential和Functional接口能让开发者快速搭出常见网络。1.3 两者并不是非此即彼从能力边界来说PyTorch 同样可以做部署TensorFlow 同样适合研究和教学。很多大型项目实际上是混合使用模型在 PyTorch 里训练和验证训练完成后导出为 ONNX 格式再接入 TensorRT 或 TensorFlow 推理链路。理解两套框架的 API 设计反而能让你更清楚地判断某个功能到底应该在哪套生态里完成。2. 环境准备与版本说明2.1 安装前的整体规划环境冲突是深度学习新手最常见的问题。PyTorch 和 TensorFlow 依赖的 CUDA 组件、底层库并不完全一致如果直接装进同一个 Python 环境很容易出现“装好这个、另一个就无法 import”的情况。最稳妥的做法是使用 Anaconda 分别创建两个独立环境pytorch_env - Python 3.10 PyTorch 2.x torchvision tf_env - Python 3.10 TensorFlow 2.x隔离之后两个环境互不干扰即使一个环境被装坏也不影响另一个。本文示例基于 Windows 10/11 或 Ubuntu 20.04/22.04Python 3.10。你的机器 CUDA 版本可能不同安装命令会有差异但排查思路是通用的。2.2 PyTorch 环境搭建先创建环境conda create -n pytorch_env python3.10 -y conda activate pytorch_env接着检查 GPU 驱动信息nvidia-smi上半部分输出里的CUDA Version: 12.x指的是驱动能够支持的最高 CUDA 版本安装 PyTorch 时可以选择等于或低于它的 CUDA 运行时版本。例如驱动支持 CUDA 12.1你可以安装 cu121也可以选更保守的 cu118。PyTorch 官网会根据你的系统环境生成安装命令。以 cu121 为例pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121如果官方源下载较慢可以换用国内镜像源例如清华源、阿里源。CPU 版本安装更简单pip install torch torchvision torchaudio这里有一个值得注意的坑创建环境时 Python 版本不要太新。某些 PyTorch 版本尚未适配最新的 Python 小版本会导致 pip 找不到对应 wheel 包。目前 Python 3.9 到 3.11 在大多数 PyTorch 版本中都比较稳妥。2.3 TensorFlow 环境搭建TensorFlow 安装比 PyTorch 特殊一点尤其是 Windows 平台。先创建独立环境conda create -n tf_env python3.10 -y conda activate tf_envCPU 版本pip install tensorflow-cpu如果是在 Linux 上想获得 GPU 支持可以安装pip install tensorflow需要注意从 TensorFlow 2.11 开始Windows 原生 pip 不再直接提供 GPU 支持。想在 Windows 上使用 GPU 运行较新版本的 TensorFlow比如 2.16、2.18 等推荐通过 WSL2Windows Subsystem for Linux或 Docker 安装。如果你的项目必须用 Windows 原生环境那么 TensorFlow 2.10 是最后支持原生 GPU 的版本。很多人装完 TensorFlow 后一运行就报 DLL 或 GPU 相关错误基本都是这个原因。Jetson 这类 ARM 嵌入式平台要特别注意不要直接从 PyPI 安装torch大概率没有对应架构的 wheel。Jetson 上通常使用 NVIDIA 针对 JetPack 系统编译的预编译包安装前先确认 JetPack 版本与 CUDA 版本再到 NVIDIA 官方资源中找匹配的.whl。2.4 安装后的检查清单安装完成后分别进入两个环境检查。PyTorch 检查import torch print(torch.__version__) print(torch.cuda.is_available()) print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else CPU only)TensorFlow 检查import tensorflow as tf print(tf.__version__) print(tf.config.list_physical_devices(GPU))如果torch.cuda.is_available()返回False优先检查驱动版本、PyTorch 对应的 CUDA 运行时版本是否匹配以及当前虚拟环境里是否真的安装了 GPU 版本。3. 核心 API 与编程范式对比3.1 张量操作两个框架都以张量为基本数据结构只是类型名不同。PyTorch 是torch.TensorTensorFlow 是tf.Tensor底层都支持 GPU 运算和自动微分。下面分别创建两个张量# PyTorch import torch a torch.tensor([1.0, 2.0, 3.0]) b a * 2 print(b)# TensorFlow import tensorflow as tf a tf.constant([1.0, 2.0, 3.0]) b a * 2 print(b)在 PyTorch 中默认设备是 CPU要用 GPU 必须手动.to(cuda)TensorFlow 检测到 GPU 后会自动尝试使用 GPU。两者的哲学差异在这里也能看出来PyTorch 把设备管理交给开发者更显式TensorFlow 更自动但也容易让人忽略资源分配细节。3.2 模型定义PyTorch 推荐继承torch.nn.Module并在__init__中定义子层在forward中定义前向传播import torch.nn as nn class MyNet(nn.Module): def __init__(self): super().__init__() self.fc nn.Linear(128, 10) def forward(self, x): return self.fc(x)TensorFlow 中既可以写函数式模型也可以继承tf.keras.Modelimport tensorflow as tf from tensorflow.keras import layers, models # 方式一Sequential model models.Sequential([ layers.Dense(128, activationrelu), layers.Dense(10, activationsoftmax) ]) # 方式二子类化 class MyNet(tf.keras.Model): def __init__(self): super().__init__() self.fc1 layers.Dense(128, activationrelu) self.fc2 layers.Dense(10, activationsoftmax) def call(self, x): x self.fc1(x) return self.fc2(x)子类化时TensorFlow 的call方法对应了 PyTorch 的forward。调用方式都是model(x)只是内部逻辑分别走forward或call。3.3 训练循环PyTorch 的训练循环通常手动编写非常透明optimizer.zero_grad() outputs model(inputs) loss loss_fn(outputs, labels) loss.backward() optimizer.step()TensorFlow 2 使用tf.GradientTape()显式记录求导过程with tf.GradientTape() as tape: predictions model(inputs) loss loss_fn(labels, predictions) gradients tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(gradients, model.trainable_variables))两者本质相同前向传播、计算损失、反向传播、更新权重。区别在于 PyTorch 把求梯度隐藏在backward()内部TensorFlow 则通过tape.gradient()把“求梯度”这一步显式交给你。3.4 数据集与数据管道PyTorch 的数据处理通常基于torch.utils.data.Dataset和DataLoader你可以自定义 Dataset 类用迭代器批量取数。TensorFlow 推荐使用tf.data.Dataset它能把文件读取、预处理、乱序、批处理、预取组合成一条完整的数据管道。在中小型项目里两套体系差别不大。但在大数据量或生产链路中tf.data提供了更丰富的并行预处理能力PyTorch 则更强调和 Python 生态的无缝衔接方便做复杂的在线数据增强。3.5 PyTorch 2.6 权重加载变化从 PyTorch 2.6 开始torch.load的weights_only参数默认值变成了True。这个改动是为了提升安全性避免恶意 pickle 文件在反序列化时执行任意代码。带来的直接影响是一些旧代码直接torch.load(model.pth)会报WeightsUnpickler相关异常。如果你加载的是完全可信的模型权重可以显式设置weights_onlyFalsestate_dict torch.load(model.pth, map_locationcpu, weights_onlyFalse)如果加载的是标准state_dict默认的weights_onlyTrue通常能够直接工作。遇到这类报错时先判断模型文件来源再决定是否关闭限制。4. 完整实战——手写数字识别两种实现这里用 MNIST 手写数字识别做例子。原因是它对算力要求很低两套框架都内置了数据下载接口加上代码量不大正好可以同题对比。4.1 项目结构dl-stack/ ├── mnist_pytorch.py ├── mnist_tensorflow.py └── README.md两个脚本互相独立分别在自己的 conda 环境中运行。4.2 PyTorch 版 CNN 完整代码文件路径mnist_pytorch.pyimport 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,)) ]) train_dataset datasets.MNIST(root./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(root./data, trainFalse, downloadTrue, transformtransform) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size256, shuffleFalse) # 2. 定义 CNN 模型 class CNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 32, kernel_size3, padding1) self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) self.pool nn.MaxPool2d(2, 2) self.relu nn.ReLU() self.fc1 nn.Linear(64 * 7 * 7, 128) self.fc2 nn.Linear(128, 10) def forward(self, x): x self.pool(self.relu(self.conv1(x))) # 28 - 14 x self.pool(self.relu(self.conv2(x))) # 14 - 7 x x.view(-1, 64 * 7 * 7) x self.relu(self.fc1(x)) return self.fc2(x) model CNN() device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 3. 训练 for epoch in range(5): model.train() running_loss 0.0 for images, labels in train_loader: images, labels images.to(device), labels.to