ARTICLE DETAIL

资讯详情

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

深度学习框架怎么选:PyTorch与TensorFlow环境搭建、安装与实战避坑指南

深度学习框架怎么选:PyTorch与TensorFlow环境搭建、安装与实战避坑指南 PyTorch还是TensorFlow这个问题我在技术社群里几乎每周都能看到提问的往往不是纠结哪个框架更“高级”而是刚装完Anaconda、准备迈出深度学习第一步的新手。其实我自己当年也在这个岔路口卡了很久网上教程一半说TensorFlow部署能力强另一半吹PyTorch写起来顺手两边的截图还不一样看得人更懵。这篇内容我想用一份过来人的视角把两个框架的真实差异、环境搭建、安装流程、实战代码风格以及新手最容易踩的坑一起捋清楚让刚接触深度学习的人能少走弯路。无论你是学生、转行做算法的、还是要跑AI模型的工程师这都是一份可以直接参考的上手指南。1. 这场“二选一”背后的真实格局两个框架到底各自强在哪1.1 PyTorch凭什么成为学术界的主流PyTorch在2017年发布当时TensorFlow已经在深度学习圈子里建立了很强的统治力。可PyTorch硬是靠“动态计算图”杀出了一条路。什么叫动态图一句话解释你写代码的时候模型的计算图是边运行边搭建的print中间结果、用if语句控制流程都很自然调试的时候就像调试普通Python代码一样。这对搞研究、做实验的人来说是致命的吸引力。另一个关键点是PyTorch和Python的融合度太高了。它不像某些框架那样有自己的一套“方言”你只要会写NumPy基本就能直接上手PyTorch。后来HuggingFace的Transformers库把几乎所有主流预训练模型都做成了PyTorch优先2024年的实际趋势也是如此翻开各大顶会论文开源的官方代码十有八九是PyTorch写的甚至很多本来就是PyTorch。对新手来说这意味着你搜一个报错、找一段参考代码大概率搜到的是PyTorch版本学习成本直接降低。1.2 TensorFlow的底牌工程化与部署生态TensorFlow比PyTorch早开源两年底子是Google的走的是“静态计算图”路线。早期版本里你得先构建一张完整的计算图然后塞进Session里去跑。这种写法对研究人员来说非常难受但换来的是部署时的稳定性计算图是固定的可以在服务器上高效执行可以做各种图优化。2019年TensorFlow 2.0发布之后Google把默认模式改成了eager execution动态执行Keras也正式收编为官方高层API。这代大改版让TensorFlow的入门门槛降了一大截但很多老项目留下的TensorFlow 1.x代码也成了一笔“历史包袱”。到了2024年TensorFlow在学术界的声音明显弱于PyTorch但别急着否定它Google的生态里TF Serving、TensorFlow Lite、TensorFlow.js还有TPU的配套支持依然是非常完整的部署链路。如果你的目标很明确是搞工业级落地TensorFlow依然有它的位置。1.3 一张表看明白两个框架的定位差异对比维度PyTorchTensorFlow计算图动态图调试直观默认动态执行核心优势在静态导出代码风格贴近Python/NumPy自由度高高层API封装完整Keras上手快学术生态顶会论文主流Transformers默认支持早期积累多但近年新论文占比下降工业部署主要靠ONNX、TorchScript中转TF Serving / Lite / JS全套移动端/嵌入式可用但不如TFLite成熟TFLite非常成熟新手学习成本需要自己写更多“细节”高层API帮你省掉不少细节看到这里你应该明白了这不是“谁比谁强”的问题而是“你的目标场景更适合谁”的问题。接下来就是选型部分。2. 做决定之前先想清楚你要用框架解决什么问题2.1 按真实用途对号入座我见过太多新手犯同一个错误先跟风选了框架学了一个月发现方向不对再换框架重头学。为了减少这种无效消耗建议按下面几种典型场景来选学生、科研人员、想复现论文直接上PyTorch。论文代码、预训练权重、社区讨论几乎都围绕PyTorch展开你用PyTorch去复现一篇论文遇到的障碍最少。目标明确要进企业做模型部署先了解公司的技术栈。有些公司内部大量沉淀的是TensorFlow生态那你就老老实实学TensorFlow如果公司是互联网大厂里偏新业务线的PyTorch ONNX的路径更常见。打比赛、做个人项目、快速验证想法PyTorch。它的灵活性让你在改模型、调试时效率高很多。要用到现成的端侧或Web端能力身边没有专门做部署的同事TensorFlow的生态更省心TFLite和TensorFlow.js的例子极多踩坑资料也全。坦白说大部分新手都属于前三类。所以我个人一贯的建议是默认先学PyTorch。理由很简单它现在就是深度学习社区里的“通用语言”。2.2 框架只是工具不要陷入“反复横跳”还有一类新手的问题不是不知道该选什么而是学了A之后听说B更好于是又去学B再后来又看到C陷入“框架收集癖”。我自己带过的实习生里真有简历上写着精通五个框架、实际写模型抓耳挠腮的。深度学习入门最忌讳的就是把精力消耗在框架切换上。算法、模型结构、数据怎么处理、训练怎么调参这些才是真正值钱的能力。框架层的东西语言是相通的卷积在哪家都是卷积LSTM在哪家都是LSTM。选定一个框架把它用透哪怕未来需要换迁移成本远比你想象的低。3. 环境搭建Python、CUDA、cuDNN、conda的版本“占位赛”3.1 先搞懂几个名词到底在说什么新手一看安装教程就头大先装Python再装Anaconda还要看CUDA版本、cuDNN版本。这些名词其实不复杂Anaconda一个Python发行版自带conda包管理工具。conda能帮我们创建独立的虚拟环境给每个项目配一套专属的Python和包互不干扰。这就是为什么教程里几乎都会让你装Anaconda。CUDANVIDIA显卡的并行计算平台。装CUDA相当于给显卡装“驱动之上的驱动”深度学习框架要调用GPU就靠它。cuDNN基于CUDA的深度神经网络加速库可以理解为专门为神经网络计算做优化的“加速补丁”。PyTorch/TensorFlow版本框架本身对应着不同的CUDA编译版本装了不匹配的版本就会出现gpu无法使用、或者直接报“not supported device”。一句话牢记你要保证的是“显卡驱动 ≥ CUDA版本 ≥ PyTorch/TensorFlow编译时的CUDA版本”。比如PyTorch官网写的是cu118意思就是它基于CUDA 11.8编译你机器上驱动支持的CUDA版本必须不低于11.8。3.2 Windows、WSL、Ubuntu、CentOS下的环境准备差异环境准备这块不同系统的坑点差别很大我分别说一下Windows系统最常见的做法是装Anaconda然后用conda创建环境。自带的显卡驱动一般已经在设备管理器里装好只要去NVIDIA官网确认驱动够新就行。PyTorch和TensorFlow都有Windows官方支持日常学习完全够用。唯一的建议是路径不要带中文和空格这是新手最容易踩的雷。WSLWindows Subsystem for Linux热点词里专门有人搜“pytorch环境搭建wsl”说明很多人已经知道WSL好用了。WSL2确实适合做深度学习实验因为它在Windows里给你一个原生的Linux环境行为和Linux服务器一致代码不容易出现“Windows能跑Linux不能跑”的尴尬。在WSL里装PyTorch注意两点第一WSL2默认支持GPU直通Windows侧装好最新的NVIDIA驱动就够了第二WSL里用的是Linux版CUDA工具包官网命令直接选Linux Tab即可不需要在WSL内再装一遍驱动。Ubuntu系统服务器深度学习最常见的组合是Ubuntu Anaconda 官方安装命令。安装前先执行nvidia-smi看驱动版本。新手常犯的错是跳到NVIDIA官网下CUDA安装包其实如果不用编译CUDA源码你只需要一个满足要求的驱动PyTorch的conda和pip包会把CUDA运行时一起带上不用单独装完整CUDA Toolkit。CentOS 7老服务器上很常见的热搜词。CentOS 7的系统库比较旧装Anaconda时如果遇到缺少libXrender之类的提示yum install对应依赖就行。Python版本选择上CentOS 7自带的Python很老务必先用Anaconda建新环境不要动系统Python不然系统的yum工具都可能被弄坏。链接界面比如你下载官方Anaconda装好后运行conda create -n pt python3.10再按正常命令安装PyTorch即可。3.3 Python版本与PyTorch版本的对应关系PyTorch对Python版本的支持是渐进的新版PyTorch会逐渐淘汰老版本Python。我给一个通用对照表方便你做环境时参考PyTorch版本建议Python版本备注1.8 ~ 1.10Python 3.6 ~ 3.9老版本项目常见1.10最好用3.81.11 ~ 1.13Python 3.7 ~ 3.101.11适合部分老显卡注意CUDA兼容2.0 ~ 2.1Python 3.8 ~ 3.11兼容性较好2.2 ~ 2.4Python 3.8 ~ 3.122024年最新系列推荐3.10/3.11TensorFlow的对应规则也类似2.10及以前的版本对Windows GPU支持很好2.11以后官方不再提供Windows原生GPU包这就导致很多人在Windows上装TensorFlow GPU版时只能退回2.10。这个问题我一会儿讲安装时再细说。4. 安装实操PyTorch和TensorFlow从零装到能跑4.1 用conda创建独立环境别把所有东西装进base无论是PyTorch还是TensorFlow我的第一建议永远是新建一个专用conda环境不要图省事装进base。原因很简单你以后的项目会依赖不同的Python版本不同框架之间还可能互相冲突共用一套环境迟早会出问题。# 创建并激活环境名字随意我用py310表示Python3.10 conda create -n py310 python3.10 -y conda activate py310激活后可以用python --version确认Python版本。这个环境就是你的“干净实验区”。后面所有安装命令都在这环境里执行。4.2 PyTorch安装命令解析CUDA版本到底怎么选PyTorch官网会根据你的操作系统、包管理器、CUDA版本生成安装命令。对于2024年来说最省心的安装方式是直接使用官网命令。先打开官网pytorch.org选好你的配置它会给你类似这样的命令# 示例Windows pip CUDA 11.8 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 示例Linux conda CUDA 11.8 conda install pytorch torchvision torchaudio pytorch-cuda11.8 -c pytorch -c nvidia这里的关键点在于CUDA版本选择。我的建议很简单先用nvidia-smi查看驱动支持的最高CUDA版本。选一个不高于这个版本的官方CUDA包。比如驱动显示CUDA最高支持12.4那装cu121或cu118都行通常选较新的cu121。不要因为显卡是老型号就保守用1.x老版本。老显卡一般也可以装合适的PyTorch 2.x除非显卡计算能力太老比如Maxwell架构之前的才需要考虑PyTorch 1.11这类兼容版本。装完验证是否成功这是最重要的步骤python -c import torch; print(torch.__version__); print(torch.cuda.is_available())如果输出True说明GPU已经可以被PyTorch调用了。这一步要是输出False先别慌下面第7节专门讲排查。国内下载慢的问题可以在pip命令里加镜像源比如清华源pip3 install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple注意加了镜像源后默认装的是CPU版因为PyPI源里放的是CPU轮子想要GPU版依然需要指定--index-url https://download.pytorch.org/whl/cu118。这是很多新手没注意到的坑。4.3 TensorFlow安装CPU版与GPU版的选择TensorFlow的安装比PyTorch微妙一点。主要因为官方从2.11开始不再维护Windows原生GPU包所以如果你在Windows上需要GPU加速绝大多数人退到2.10# Windows GPU推荐直接装2.10 pip install tensorflow2.10 # Linux GPU pip install tensorflow # CPU版本任何平台可以先拿来练手 pip install tensorflow-cpuLinux上装了2.11以上版本的话GPU支持是内置的但需要你在系统层面装好CUDA和cuDNN装起来比PyTorch的“开箱即用”繁琐一点。因此很多新手在Linux上更愿意用PyTorch起步。验证TensorFlow是否正常python -c import tensorflow as tf; print(tf.__version__); print(tf.config.list_physical_devices(GPU))如果能打印出GPU设备列表就说明安装成功。如果只打印了CPU设备说明TensorFlow没找到CUDA环境大概率是版本匹配问题。5. 用同一个MNIST任务实测两个框架代码风格和上手体验5.1 PyTorch版动态图、自定义训练循环理论说再多不如跑一个代码。我用MNIST手写数字识别来对比这个任务相当于深度学习的“Hello World”。PyTorch的典型代码结构分成几步定义网络结构、定义损失函数、定义优化器、写训练循环。关键是你可以用标准的Python控制流来写循环每一步都能打断点查看张量数值import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms # 数据加载 transform transforms.ToTensor() train_data datasets.MNIST(root./data, trainTrue, downloadTrue, transformtransform) train_loader torch.utils.data.DataLoader(train_data, batch_size64, shuffleTrue) # 定义网络 class Net(nn.Module): def __init__(self): super().__init__() self.fc nn.Linear(28*28, 10) def forward(self, x): x x.view(x.size(0), -1) return self.fc(x) model Net() criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 训练循环 for epoch in range(3): for images, labels in train_loader: outputs model(images) loss criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() print(fepoch {epoch1}, loss: {loss.item():.4f})PyTorch的训练逻辑非常直白前向算出loss反向传播更新参数全是显式的。对于想搞懂深度学习原理的新手这种显式流程其实是很好的学习材料你能清楚地看到每一步发生了什么。5.2 TensorFlow版本Keras的高层API更“省心”TensorFlow用Keras高层API写同一个模型代码更短封装更厚import tensorflow as tf # 数据加载 (x_train, y_train), _ tf.keras.datasets.mnist.load_data() x_train x_train / 255.0 # 定义与训练一体 model tf.keras.Sequential([ tf.keras.layers.Flatten(input_shape(28, 28)), tf.keras.layers.Dense(10, activationsoftmax) ]) model.compile(optimizeradam, losssparse_categorical_crossentropy, metrics[accuracy]) model.fit(x_train, y_train, epochs3, batch_size64)你看model.fit一行就把PyTorch里那个手动训练循环干完了。Keras的抽象层替你做掉了训练细节对只想快速出结果的新手来说确实友好。但代价是你不太容易搞清楚内部发生了什么一旦需要自定义训练循环或者要写一些很特殊的模型结构Keras高层API的限制感就出来了。5.3 两者对比新手在真实项目中会感受到什么从上面的代码你肯定能感受到两种风格。我再举一个场景你在调试模型时发现中间某个张量维度不对想打印一下它的shape。在PyTorch里直接print(x.shape)写在任意一行都行因为动态图每次运行都会执行到那里。而一旦你进入复杂的自定义模型这种“随时可以加print调试”的能力会帮你节省大量时间。反过来如果你要快速搭建一个标准模型、跑一个简单任务TensorFlow的Keras确实更省力。这就回到了第2节说的选型问题目的是学习原理、写灵活模型、复现论文PyTorch更合适目的是快速搭建标准流程、有现成部署链路TensorFlow很省心。再提一个热门搜索词里的具体场景用PyTorch在seq2seq模型里写一个通用的attention模块。这个场景特别能体现PyTorch的优势。你可以把attention的计算过程拆成一个普通的Python类比如计算score、再做softmax、再加权求和每一步都是显式的张量运算调试和修改都极其自然。同样的逻辑在静态图框架里会绕很多。这也是为什么NLP研究和实战社区里PyTorch几乎成了默认选项。6. 从模型到产品PyTorch转ONNX和TensorFlow部署思路6.1 为什么现在“训练用PyTorch、部署走ONNX”这么常见新手学到一定程度就会碰到部署问题。你训练好的PyTorch模型要给别人用别人不一定会装PyTorch就算装了不同的PyTorch版本也可能导致加载失败。所以行业里出现了一条越来越主流的路用PyTorch训练模型然后导成ONNXOpen Neural Network Exchange中间格式再交付给不同的推理引擎。ONNX像是一个“标准翻译器”它把PyTorch或者TensorFlow的模型翻译成一套通用的计算图描述。下游可以用ONNX Runtime、TensorRT、OpenVINO等引擎去加速推理也可以方便地转到其他设备。对新手来说ONNX Runtime是最好上手的推理引擎pip安装后直接用即可不需要依赖PyTorch。6.2 PyTorch转ONNX的实操步骤与动态轴坑点PyTorch导出ONNX用的是torch.onnx.export代码本身不长坑主要在参数细节上import torch # 假设model是你的训练好的模型device是cpu或cuda model.eval() dummy_input torch.randn(1, 3, 224, 224) # 一个假输入仅仅为了确定输入形状 torch.onnx.export( model, dummy_input, model.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}}, opset_version11 )这里三个关键点dummy_input的形状必须和真实输入一致一般用1张图的尺寸Batch维度设为1。dynamic_axes指定动态轴。不加这个参数的话你导出的模型会把batch size固定死成1推理时不能改变。最常见的用法就是把第0维batch给成动态这样推理时一次能处理张数可变的输入。opset_version是ONNX算子集的版本默认值一般也行但如果你用的是比较新的算子可能要升到13或17。有些老旧推理引擎对太新的opset支持不好需要根据实际部署环境试验。导出后一定要验证import onnxruntime as ort import numpy as np sess ort.InferenceSession(model.onnx) input_name sess.get_inputs()[0].name result sess.run(None, {input_name: np.random.randn(1, 3, 224, 224).astype(np.float32)})[0] print(result.shape)如果这一步报错最常见的原因是模型里用了某个ONNX不支持的算子或者某个动态轴的配置不对。把报错关键词复制到搜索引擎基本都能找到解决方案。6.3 TensorFlow的SavedModel与TF Serving部署路径TensorFlow的官方交付格式是SavedModel用Keras训练好的模型一行代码就能导出model.save(saved_model/my_model)之后可以用tensorflow/serving镜像起一个推理服务或者用tflite_convert转成TFLite跑在移动端。当你需要管理多个模型版本、做线上A/B测试、用标准REST/gRPC接口提供服务时TF Serving是真心好用。这方面TensorFlow的成熟度确实比PyTorch高一大截。如果你最终要落地到Android端的应用TFLite的生态最完善资料最多。这一点要客观承认TensorFlow的生命力很大程度上也来自部署环节。7. 新手最容易踩的坑CUDA版本、设备不支持、显存不足7.1 “PyTorch不支持设备”到底是怎么一回事“绘世启动器显示pytorch不支持设备”是我看到的热搜词里一个非常有代表性的场景。很多人在用Stable Diffusion绘世整合包时打开启动器弹出一行提示说PyTorch不支持当前设备。这个问题的本质就是显卡型号太老或者版本组合不兼容。绘世这类整合包是内置了PyTorch和CUDA环境的它检测到当前显卡计算能力不满足要求时就会报警。常见原因有几个显卡是NVIDIA太老的型号比如GTX 700系列及更早期计算能力太低新版PyTorch就不支持了。解决方法是找配套的PyTorch 1.11或旧版本或者干脆换卡。显卡驱动太老导致CUDA版本识别不到。去NVIDIA官网更新驱动就能解决。显卡是NVIDIA的低端卡显存很小虽然能装上但跑起来容易OOM。遇到这类提示先确认你的显卡型号和驱动版本再决定是老实用CPU跑、降级PyTorch版本还是换平台。不要一上来就怀疑整合包坏了。7.2 CUDA版本不匹配的排查流程安装完PyTorch后如果torch.cuda.is_available()返回False按这个顺序排查打开命令行输入nvidia-smi看显卡是否被识别、驱动版本是多少。如果这个命令都找不到显卡先解决驱动问题。在Python环境里输入import torch; print(torch.version.cuda)看你安装的PyTorch自带的是哪个CUDA版本。把这个版本和nvidia-smi里显示的驱动支持的最高CUDA版本对比前者不能高于后者。检查conda环境里是不是有多个CUDA包混在一起。如果之前瞎试装了很多最干净的做法是重开一个conda环境从头装一遍。如果你是在WSL里用的确认Windows侧的驱动足够新且在WSL里没有再去装一个不匹配的NVIDIA驱动。很多时候新手折腾半天最后发现就是环境太脏了重开环境解决90%的问题。7.3 显存、CPU版本训练慢、图像界面卡死等常见问题还有几个频率很高的坑我直接列出来显存不足CUDA out of memory把batch_size调小这是最快的方法。如果已经很小的还是爆检查是不是模型太大考虑换一个精简版模型。有人以为是显存泄漏其实绝大多数新手就是batch设大了。训练慢到怀疑人生检查你是不是装的CPU版PyTorch。很多人下载了PyTorch默认包这个默认是CPU版本GPU根本没参与训练。验证方法就是torch.cuda.is_available()返回True才算GPU生效。下载慢、装到一半失败用国内镜像源或者避峰下载。PyTorch的whl包很大一个G很常见网络不稳就换个时间段。conda源也可以换成清华源或中科大源改一下.condarc配置就行。多个人共用的服务器上乱装环境养成用conda建独立环境的习惯不要看到什么顺手就pip install到base。8. 我的实际体会给新手的一套启动方案讲了这么多最后分享一点我自己的经验。如果你今天刚装了Anaconda正在面对选择困难我给一套可以直接照抄的启动方案第一步装PyTorch。理由前面说过它在当前生态里资料最多、踩坑答案最好搜、对新手最宽容。你可以完全不装TensorFlow先花一个月把PyTorch的基础跑熟。第二步把环境装到能用就行不要追求完美。比如CUDA版本我见过不少新手为了“到底装哪个CUDA版本”纠结一整天。其实你只要用nvidia-smi看一眼然后选一个官网提供的、不高于驱动版本的CUDA包就够了。就算后来发现选错了重开conda环境也就几分钟的事别怕。第三步务必跑通一个小项目。不用贪大把第5节的MNIST代码跑起来看到loss在下降、模型能出预测结果这比看十个教程都管用。跑通之后再加一个小功能比如把模型导出成ONNX用ONNX Runtime推理一下你会瞬间理解整个训练到部署的闭环。第四步遇到问题学会搜报错。把报错信息原样粘贴到搜索引擎比你在群里描述“我的模型跑不了”有用得多。PyTorch和TensorFlow的全网讨论量非常大你踩的坑99%都有前人记录。第五步等你有了一定基础再去接触TensorFlow。比如你想了解高层API怎么让训练代码变短或者想看Serving怎么部署模型再去装TensorFlow不迟。两个框架都懂一些当然好但那应该是进阶阶段的事而不是入门第一周的事。我在实际带人的过程中发现凡是能连续一周每天花两小时跑代码的人无论选哪个框架最后都能学会凡是把时间花在反复比较、反复换环境、反复纠结“我是不是选错了”的人一个月后还在原地。深度学习的门槛不在框架而在你愿不愿意早点动手跑起来。先让代码跑起来剩下的问题都会在实战中一个一个解决。
返回列表