多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

PyTorch与TensorFlow双框架工程实践:从环境搭建到ONNX部署

PyTorch与TensorFlow双框架工程实践:从环境搭建到ONNX部署 做了几年深度学习落地周围一直有人问我到底该主攻PyTorch还是TensorFlow。每次我都说这个问题本身问错了。真正该问的是这两个框架各自的边界在哪里以及你手里的项目更适合哪一个。过去我也站过队后来被现实反复教育现在可以明确告诉你——PyTorch和TensorFlow之间根本没有必要分出你死我活把两个框架都纳入自己的工具箱才是应对这个领域最稳妥的方式。这篇文章我会从工程落地的角度把如何同时掌握PyTorch和TensorFlow这件事拆开来讲包括双框架环境搭配、CUDA与Python版本对应关系、核心API的相互翻译、PyTorch模型转ONNX再走向TensorFlow Serve的完整链路还有我在实际项目里踩过的一些坑。不管你是刚入门的初学者还是被项目逼着切换框架的老兵这些内容应该都能帮你省下不少折腾时间。1. 全都要不是贪多而是工程现实1.1 两个框架各自主场的真实分布先说学术和科研这一侧。Transformer相关的论文、Vision Transformer的实现、各种新出的注意力机制我先在GitHub上搜到的复现代码大概率是PyTorch写的。HuggingFace的Transformers库主力后端就是PyTorch虽然也兼容TensorFlow但很多新模型的示例都是PyTorch优先。你在arXiv上拿到了开源权重多半也是PyTorch的.pt或者.bin文件。所以做研究、跑新模型、复现论文PyTorch基本是绕不开的。再看工业部署这侧。TensorFlow Serving在模型上线方面足够成熟TFRecord的数据管线在超大流量场景下很能打TensorRT和TensorFlow的衔接链路也被验证过很多年。很多老牌公司的算法部门线上推理还跑着TensorFlow的SavedModel。这时候如果只会PyTorch你拿到的模型权重可能都要靠同事先转一道再给你。反过来如果你只懂TensorFlow看到最新的PyTorch权重想做微调第一件事就得先转格式搞不好还被算子支持列表卡住。我自己的经验是这两个框架并不冲突它们像是同一个手艺的两种表达方式。PyTorch的即时执行让调试很舒服TensorFlow 2.x把keras和eager execution拾起来之后上手体验也没那么冷冰冰了。只有两套都顺手才能在不同场景里不慌。1.2 把两选一变成多一张底牌很多人纠结框架本质是怕学习成本太高。这个担心有道理但也没有想的那么夸张。两个框架的核心思想其实共用一套——张量、自动求导、层模块、优化器、损失函数。你在PyTorch里理解了backward()和计算图到了TensorFlow里不过是换成了GradientTape你在TensorFlow里习惯了model.fit(),到了PyTorch里用自定义训练循环多写几行反而把反向传播的流程看更清楚了。可以说掌握一个框架之后第二个框架的学习曲线会平缓很多。我见过不少从Keras上手的人再回来学PyTorch大概一周就能跑通自己的项目。反过来从PyTorch转TensorFlow也不难重点是忘掉必须用model.fit才能训练的惯性。框架只是工具真正值钱的是你对模型结构、数据流和训练过程的深层理解这部分能力在两个框架间完全通用。2. 双框架环境搭建CUDA版本、Python版本和包管理器的三角关系2.1 环境隔离是第一原则要让PyTorch和TensorFlow和平共处第一原则就是别把它们装进同一个Python环境。虽然理论上可以共存但是它们对CUDA相关库的依赖容易互相拉扯你升级TensorFlow时顺手覆盖了某个动态库PyTorch可能就开始报CUDA错误排查起来非常痛苦。我见过太多次这种问题后来学乖了一律用环境隔离。方案上我推荐两种。一种是Anaconda创建不同虚拟环境比如conda create -n torch python3.10和conda create -n tf python3.10。用的时候conda activate torch或conda activate tf干净利落。另一种是装WSLWindows Subsystem for Linux在Linux子系统里面做深度学习开发宿主机保持日常使用。WSL对GPU的支持现在已经很完善了CUDA在WSL里的体验很接近原生Linux很多在Windows上折腾cuda和pytorch的人换到WSL后瞬间顺了。提示环境隔离要趁早。项目到中期再拆分环境会把所有依赖都绑在一起改一个包版本就可能引爆另一个问题。2.2 一步都省不了的CUDA与PyTorch版本核对很多人装PyTorch失败不是网络问题是版本不匹配。PyTorch、CUDA和Python版本之间有一张对应关系表CUDA Toolkit本身还可能和显卡驱动不兼容。我的建议是先看显卡驱动支持的CUDA版本再根据PyTorch的官方说明安装对应的PyTorch版本。具体步骤通常是这样的在终端输入nvidia-smi看右上角的CUDA Version那是当前驱动支持的最高CUDA版本。这个值只需要大于等于你要安装的CUDA runtime版本即可。打开PyTorch官网的安装页面选择你的操作系统、包管理器和CUDA版本复制生成的命令。如果你想要Pytorch 1.11这种历史版本可以进PyTorch官方历史版本列表里复制对应的pip install torch1.11.0cu113 ...命令。不要急着装最新版。如果你的显卡是老卡新版PyTorch可能已经放弃支持了。这时候选择旧版的cu111或者cu102反而更稳妥。一个很常见的报错是CUDA unavailable或者PyTorch not compiled with GPU support十有八九是装成了CPU版本的PyTorch。Windows用户尤其要注意pip install torch默认装的是什么版本要确认好最好直接把cu118这样的后缀带进URL里不要依赖默认源。2.3 Ubuntu、WSL和CentOS上的安装要点Ubuntu是深度学习最顺的系统。一般流程是先装NVIDIA驱动再装CUDA Toolkit和cuDNN最后用conda或者pip建环境。需要注意的坑是PyTorch官方默认的pip源在国内很慢换国内镜像源记得要同步配conda或者pip的配置否则下到一半失败的概率很高。WSL的安装思路则是另一条路。你在Windows侧装好NVIDIA驱动然后WSL里不用再装驱动直接在WSL内安装CUDA Toolkit和PyTorch就行。很多人在pytorch环境搭建wsl上卡住多半是因为把Windows侧的CUDA和WSL侧搞混了。简单说WSL里的GPU访问走的是Windows驱动层的映射你只需要确保驱动版本足够新然后在WSL里用nvidia-smi能看到显卡再按Linux方式装PyTorch即可。CentOS这种服务器系统也有不少人用。CentOS最大的问题是系统自带的Python版本很旧很多依赖包用yum装不上。常规做法是先装Miniconda或Anaconda到用户目录然后在conda里创建Python 3.8/3.9/3.10环境再安装PyTorch或TensorFlow。这个过程里最常见的坑是glibc版本太低导致新版Python启动失败老系统上需要先升级系统库或者选用更早期的Python版本。3. 两套框架的同步翻译能力从模型定义到训练循环3.1 核心API对照表真正把两个框架都用起来之后你会发现80%的代码都能在另一个框架里找到对应物。下面这张表是我自己整理的很常用功能PyTorchTensorFlow张量对象torch.Tensortf.Tensor定义模型层torch.nn.Moduletf.keras.Model / tf.keras.layers.Layer全连接层torch.nn.Lineartf.keras.layers.Dense卷积层torch.nn.Conv2dtf.keras.layers.Conv2DReLU激活torch.nn.ReLUtf.keras.activations.relu / tf.nn.relu损失函数torch.nn.CrossEntropyLosstf.keras.losses.CategoricalCrossentropy / SparseCategoricalCrossentropy优化器torch.optim.Adamtf.keras.optimizers.Adam加载数据集torch.utils.data.DataLoadertf.data.Dataset前向传播model(x)model(x, trainingTrue)反向传播loss.backward() optimizer.step()with tf.GradientTape() as tape: ... optimizer.apply_gradients这个表对写过Keras的人特别友好。比如你原来只写过model.fit转到PyTorch后觉得训练循环很难受其实只是还没习惯。PyTorch的训练循环本质就是for inputs, labels in dataloader: optimizer.zero_grad() outputs model(inputs) loss criterion(outputs, labels) loss.backward() optimizer.step()而在TensorFlow 2.x里对应的写法是for step, (inputs, labels) in enumerate(dataset): with tf.GradientTape() as tape: outputs model(inputs, trainingTrue) loss criterion(labels, outputs) grads tape.gradient(loss, model.trainable_variables) optimizer.apply_gradients(zip(grads, model.trainable_variables))对比看下来逻辑完全一样无非是谁在管理梯度的差异。理解了这一层你就不再需要怕切换框架了。3.2 静态图与动态图别被概念吓住PyTorch默认是动态图也就是说前向传播每次执行都会重新构建计算图这让调试非常直观——你可以print网络中的任意张量可以用Python的if和for来写模型逻辑。TensorFlow 2.x默认也是eager执行看起来和PyTorch的行为一致但它保留了静态图能力比如tf.function把Python函数编译成计算图提升性能。实际开发中动态图适合研究和快速原型因为改起来没有编译延迟静态图适合反复执行同一个模型比如线上serving图优化能带来明显提速。明白了这个区别你就能解释为什么很多PyTorch项目在部署时要转成ONNX或者TorchScript——其实就是把动态图定成一个静态计算流程。理解层面的差异比API差异更重要。动态图让你调试方便静态图让你上线安心两个框架恰好代表了这两种思路而已。4. 桥接关键路径PyTorch模型转ONNX再走向TensorFlow生态4.1 为什么要以ONNX作为中转站很多团队的实际状态是研究端用PyTorch快速迭代生产端又希望享受TensorFlow Serving或者TensorRT的成熟体系。直接跨框架导权重是一件很麻烦的事但ONNXOpen Neural Network Exchange把这个问题变得简单的多。ONNX相当于一个公开的模型中间表示上游框架把模型导出为ONNX下游框架再把ONNX导入成自己的格式。PyTorch官方提供了torch.onnx.export把模型导出成ONNX。随后在TensorFlow侧导入ONNX最常见的是用onnx-tf这个库里的onnx_tf.backend.prepare之类的接口。我实际测试下来比较标准的CNN、Transformer、BERT类模型都可以走通。转换后如果遇到算子不支持通常可以从以下三个方向排查模型里用了比较新的算子ONNX版本没跟上升级ONNX版本或者换成一个新一点的Opset。动态shape问题。ONNX导出时的dynamic_axes参数没有配好导致转换后模型只能在固定的输入尺寸上运行。一般需要显式声明dynamic_axes{input: {0: batch}}。自定义层根本没有ONNX对应实现。这种情况需要你把自定义逻辑用标准算子重新组合或者直接通过ONNX注册自定义算子。转换完之后别急着部署一定要对比输出。把同一批输入分别跑PyTorch原模型和转换后的模型比较输出张量的数值误差。ONNX转换在数学上应该是接近一致的但浮点顺序不同会带来极小的误差如果误差明显很大说明某个算子映射有问题。4.2 DataLoader到tf.data Dataset的迁移除了模型本身数据管线也可以无缝迁移。PyTorch里习惯用torch.utils.data.Dataset定义每个样本的读取再用DataLoader加batch、shuffle和num_workers。到了TensorFlow里对应的正确姿势是tf.data.Dataset.from_tensor_slices搭配map、batch、shuffle、prefetch。几个容易踩的坑num_workers对应到TensorFlow是map(num_parallel_callstf.data.AUTOTUNE)不设置的话数据读取可能成为性能瓶颈。PyTorch的DataLoader默认在迭代时才异步预取TensorFlow则需要你在pipeline末尾加prefetch(tf.data.AUTOTUNE)否则下个batch的加载会停顿。shuffle的位置要注意。PyTorch里DataLoader自带shuffle是在每个epoch打乱索引tf.data里常见的顺序是shuffle(buffer_size)在map之后或者之前要确保打乱的是样本而不是只打乱batch。如果你有大量TFRecord数据TensorFlow侧直接用tf.data.TFRecordDataset按record读取比dataloader再去读普通文件要快一个量级。这也是很多工业项目坚持在TensorFlow侧做完整训练流程的原因——数据管线的优势体现在大数据量上。5. 实战策略长期维护双框架项目的经验清单5.1 优先用PyTorch跑新模型用TensorFlow侧做稳定服务以我个人的习惯探索阶段一定在PyTorch里进行。数据加载、模型搭建、训练调参都在一个环境里完成因为PyTorch对调试的友好度真的很高。等模型在验证集上效果稳定了再导出ONNX然后到TensorFlow侧做转换、验证和上线。这套流程的收益是两方面的研究端不需要为了TensorFlow的部署限制而牺牲灵活性部署端也不需要因为PyTorch的调试便利而放弃成熟的serving体系。关键是转换环节想清楚了加上清晰的版本管理两个框架就能各司其职。举一个具体的流程在PyTorch里定义模型结构并加载训练好的权重。用torch.onnx.export导出ONNX文件设置好opset版本和dynamic_axes。在TensorFlow环境里用onnx-tf转成SavedModel。用tf.saved_model.Load加载跑一个model(x)的推理测试与PyTorch输出做误差对比。接入TensorFlow Serving或转成TensorRT engine继续后续的延迟和吞吐性能优化。这套链路跑顺之后别人还在为了模型是pt文件还是pb文件吵来吵去你已经可以在任意一侧开展工作了。5.2 环境、版本和权重管理双框架项目活下去的三根支柱双框架项目最大的风险不在写代码而在环境漂移。我有两个很实际的建议第一每个项目一开始就把环境导成清单文件。PyTorch环境用pip freeze requirements.txt或者conda env export environment.ymlTensorFlow环境同理并且最好在README里写明CUDA版本、Python版本和显卡驱动版本。等隔了三个月再回去跑老项目你就会因为这行记录少掉很多白头发。第二模型权重格式要在项目内统一。我见过一个团队有人产出.pt文件有人产出.h5文件还有人产出ONNX文件最后衔接的时候发现根本对不上。建议在项目文档里明确约定训练和迭代阶段的产物统一用.pt对外发布和部署统一用ONNX或SavedModel。这不只是文件后缀的洁癖它直接决定上下游能不能顺畅合作。5.3 学习顺序建议先深后宽再回到深如果你现在还是新手我建议的路径是先选一个框架学到能独立完成一个项目再学习另一个框架。这个顺序听起来和两个都要矛盾但你仔细想想先深后宽是先掌握一套核心逻辑再迁移到另一套时你是在框架的相似点上快速复用而不是被两套API同时搅浑。我经常和新同事说别指望在入门阶段同时学两个框架那样你只会两头都浅。先花两个月把PyTorch或TensorFlow其中一个做到能独立训练模型到时候再花一两周看另一个框架的官方教程和常见模型实现就会轻松很多。另外一个建议是日常可以偶尔试试用PyTorch跑一个Keras教程里的模型或者用TensorFlow实现一篇PyTorch仓库里的代码。这种练习不需要是项目级的训练一个小模型在MNIST或者自己的小数据集上跑通就足够。目的不是产出成果而是让你保持对两套框架手感的热度。6. 我实际遇到过的坑和最终的框架选择心得6.1 三个值得单独记下来的问题第一个是PyTorch环境搭建好了但CUDA用不了。后来发现是装了CPU版PyTorch。解决办法是卸载后重新安装torch1.13.1cu117这类的带CUDA的版本。每次新环境都记得用torch.cuda.is_available()验证一下这个步骤不要跳过。第二个是TensorFlow和PyTorch共存时某个库的版本被另一个环境覆盖掉了。我最早贪方便两个框架装在同一个conda环境里结果某天TensorFlow升级后把cudnn的版本动了一下PyTorch直接报错。从那以后我再也不混装就算只维护一个项目我也会单独建环境。第三个是ONNX转换后推理结果和目标检测模型带的后处理逻辑不匹配。模型online的转换本身没错问题出在PyTorch里有些后处理步骤写在了模型外面比如NMS非极大值抑制导出ONNX时这些逻辑默认没有包含进去。结果转出来的模型只有前向推理的输出下游接的nms就失效了。处理方案是导出时把NMS这类逻辑显式放到模型内或者干脆在TensorFlow侧重新实现一份后处理。6.2 最终心态框架只是工具别让工具定义你的视野到了这个阶段我对PyTorch和TensorFlow的态度已经很平和了。两者背后都是深度学习的基础理论——张量计算、自动微分、梯度下降、反向传播。你在调试模型时真正用到的能力比如观察loss曲线、分析梯度异常、调整学习率策略这些在一个框架里练熟了换个框架毫无障碍。最后再分享一个小技巧无论主力框架是什么我都会让自己保持至少能把同一个模型用另一个框架复现的能力。不需要每月做一次至少半年做一次。这样当团队突然要切换技术栈或者某个模型只有另一种框架的预训练权重时你不会被卡在工具选择上而是能直接把活接下来。在深度学习这个行业里比我会哪个框架更值钱的是我了解模型和数据在计算体系中如何流动。别让框架选择限制你的可能性真正该选择的是你想解决的问题本身。
返回列表