ARTICLE DETAIL

资讯详情

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

PyTorch模型安卓端部署实战:从NCNN到TorchScript的踩坑之路

PyTorch模型安卓端部署实战:从NCNN到TorchScript的踩坑之路 标题里的Torchscipts其实是我某天凌晨改代码时打出来的错别字正确拼写是TorchScript。我故意保留它因为它准确描述了那段时间的状态——为了把一个目标检测模型塞进安卓端我从NCNN试到TFLite再试到MNN全军覆没最后靠TorchScript才跑通。这篇文章不是什么成功学分享只是一份踩坑记录写给那些正准备把已有模型部署到手机、平板上的人。如果你也天真地以为模型训练完就能直接上端那这份血泪史应该能帮你省下几个通宵。1. 明明已经训练好的模型怎么一到安卓端就成了玄学1.1 项目背景一台不联网的旧平板和一个微瑕疵检测需求先交代一下我当时在做什么。项目是一条小型产线的零件外观检测需要定位工件表面非常细微的划痕、缺角和色差区域。模型用的是YOLOv5s我在backbone里加了一个很轻量的通道注意力模块训练后检测精度还不错小目标的召回率也能接受。但客户方的使用环境非常特殊一台已经用了三四年的安卓平板车间里基本没有稳定Wi-Fi数据也不能往云端传。这意味着我不能做手机通过网络请求调用电脑上部署的模型这种方案只能在平板本地完成整个推理链路。当时我心想这还不简单PyTorch导出ONNX然后转成NCNN或者TFLite在Android Studio里接个SDK不就完事了。现在回头看这份自信就是灾难的开始。1.2 我以为的部署链路和实际上的部署链路我以为的链路是这样的训练好的 PyTorch 模型 → 导出 ONNX → 转成 NCNN/TFLite/MNN 格式 → 在 Android 端加载 → 预处理图片 → 推理 → 后处理画框任何一个环节都不复杂对吧实际上真正的链路是PyTorch 模型 → 导出 ONNX 时踩算子兼容性坑 → 用转换工具转 NCNN/TFLite/MNN 时再踩一层算子坑 → 编译安卓工程时踩 SDK/NDK 版本坑 → 加载模型后要么直接崩要么输出形状不对 → 要么终于跑出来一堆框但位置完全对不上 → 回头检查预处理、归一化、后处理发现又是另一套坑用生活里的例子打比方假设你有一篇用简体中文写好的文章想拿到一个只认繁体中文的论坛发布。你以为装个翻译插件就行结果发现文章里有大量网络用语和专有名词翻译软件根本不认识翻出来驴唇不对马嘴。如果你最后找到一份同作者原始稿件直接发布那才是真正省事的路——TorchScript解决的就是这个问题它保留了PyTorch自身的语言习惯不需要二次翻译。我当时最大的认知错误就是把模型转换想得太轻量。实际上PyTorch模型和NCNN/TFLite/MNN之间的差异等同于两套完全不同的执行引擎每一个不认识的算子都可能让转换链路的某一环直接卡死。2. NCNN我栽在官方支持三个字上2.1 为什么第一站选NCNNNCNN是腾讯开源的高性能神经网络前向计算框架专为移动端优化模型体积小、推理速度快社区活跃度也高。更关键的是我搜了一圈发现很多开发者都用它来部署YOLOv5当时NCNN官方仓库里甚至直接就有YOLOv5的示例代码。这意味着什么意味着只要我老老实实用标准的YOLOv5s大概率是能走通一条已经被前人踩平的路线。可问题就出在老老实实这四个字上——我的模型不是标准YOLOv5s我在backbone里加了通道注意力模块计算图中会出现一些额外的Multiply、GlobalAveragePool、Sigmoid组合。我觉得这些都很基础NCNN肯定支持但实际验证下来发现事情没那么简单。2.2 从ONNX到NCNN的报错全记录我的操作步骤很常规先用torch.onnx.export导出ONNX然后在电脑上用onnx2ncnn做转换。第一次转换就报了一串错误我到现在都记得那种感觉——红色日志一行行往上刷其中最关键的一句大意是Can not create layer StridedSlice (type: StridedSlice)为什么会有StridedSlice因为YOLOv5早期版本的数据增强和输入处理里用了切片操作Focus结构在ONNX里会被表达成StridedSlice。NCNN当时对它的支持非常有限遇到这种层直接跳过导致后续的卷积全部对不上形状。我试过升级NCNN版本、换Vulkan版本、用ncnnoptimize工具做优化最后报错的层变了但新的错误又冒出来比如某些Resize参数组合不被识别。说实话如果只是遇到一个算子缺少支持我还不会立刻放弃。真正让我崩溃的是排查成本实在太高ONNX里的一个算子到NCNN里可能对应完全不相关的几个层中间还夹着形状推导错误。我花了两天去改模型结构、绕过某些特定写法但每修好一个错误就会在下一个环节遇到新的问题像打地鼠一样没完没了。2.3 不是NCNN不行而是特化支持和通用支持的差距后来我才想明白NCNN社区之所以能跑通YOLOv5是因为有开发者专门为这个模型适配了算子、写好了示例、甚至做好了后处理代码。换句话说这是特化支持。你的模型只要和官方YOLOv5结构稍微偏离一点就得不到这些预置红利的保护必须自己去面对算子兼容性矩阵。NCNN本身没有问题它在很多场景下确实又快又稳。但官方有示例不等于你的模型能被通用支持。我在那段时间最大的体会是移动端部署框架的算子覆盖率永远不可能和PyTorch这种完整训练框架对齐。你训练时可以随便用nn.Flatten、F.interpolate、torch.cat但这些操作到了转换工具里可能变成一连串难以处理的细节。如果你的模型结构相对新或者自定义那么第三方转换工具就会成为整个部署链路上最不可控的一环。3. TFLite 与 MNN没想到接近成功才是最折磨人的3.1 TFLite卡在转换链路和量化精度上既然NCNN走不通我的第二反应就是换TFLite。毕竟TFLite背靠Google文档齐全社区方案也多很多人说TFLite对移动端推理的支持更成熟。我选择了一条常见的路线先把PyTorch模型导出为ONNX再通过ONNX-TensorFlow转换成TensorFlow的pb模型最后用TFLite Converter转成.tflite文件。这一条链路光是工具版本就够你喝一壶的。我遇到的第一个问题依然是算子兼容。ONNX里一些动态尺寸相关的操作在转换成TensorFlow操作时会变成很奇怪的子图TFLite Converter往往不支持其中某些复合操作要么直接报错要么转换生成的tflite模型在推理时输出全零。我试过用tf2onnx、onnx-tf这类工具组合始终没有一个稳定方案。更折磨人的是转换成功后的模型在电脑上可以跑但一部署到安卓端就输出NaN或者一堆置信度接近0的框完全不可用。量化是另一个大坑。为了让模型在旧平板上跑得快一点我尝试了TFLite的动态范围量化。量化完成后模型确实从14MB缩到了7MB左右运行速度也上来了但精度损失让我直接傻眼——原本能检出的那些微小划痕在量化后几乎全部消失。后来又试了全整型量化需要提供代表性数据集做校准我选了500张训练样本做校准结果也只是有所改善远达不到可用标准。后来我查资料才明白小目标检测对特征图的敏感度极高量化误差会在低层特征上被放大微小缺陷这种本来就处于边缘的信号很容易就被量化噪声淹没了。3.2 MNN版本黑洞和运行一致性MNN是阿里的开源框架主打高算子覆盖率。当时我已经被NCNN和TFLite磨得没脾气了看到MNN官网说支持绝大多数常用算子心里又燃起一丝希望。可实际用下来MNN是我三个框架里踩坑最难受的一个。先用MNNConvert工具把ONNX转成MNN格式第一次确实比NCNN顺利报错少很多。但一到Android端就出问题官方给的so和jar包版本不匹配加载模型后调用forward直接崩溃换了对齐的版本CPU后端能跑但输出结果和Python端的PyTorch结果差距很大尤其是检测框的中心点坐标总是偏移。我一开始以为是自己后处理写错了反反复复检查Java代码甚至把MNN的Python接口拿来逐层Debug最后怀疑是某些算子融合导致了数值精度问题。还有一次是MNN在Vulkan后端和CPU后端跑出来的结果完全不一样——同一个模型、同一张输入图CPU能检出缺陷Vulkan下就漏检。这种跨后端的不可控性比完全不支持还让人头疼。因为完全不支持你至少知道问题在哪不一致意味着你连问题的边界都摸不着。3.3 三个框架的共同规律NCNN、TFLite、MNN我都试过了失败原因各不相同但背后有一条共同规律你的模型是在PyTorch生态里训练的而目标推理框架的执行逻辑和PyTorch之间存在一条不透明的翻译通道。这条通道里任何一个算子不被识别、任何一种形状推导不符合预期、任何一种精度模式不一致都会导致整个部署失败。那段时间我甚至开始怀疑是不是自己太菜了。后来和一位做端侧部署的朋友聊他说了一句话点醒我你的模型本来就是PyTorch的为什么非要先转成别的格式再部署直接上PyTorch Mobile不好吗我这才把目光投向TorchScript。4. TorchScript救场原生格式如何绕过所有中间层4.1 突然醒悟我为什么非要绕一圈TorchScript是PyTorch官方提供的序列化格式它把训练好的模型打包成一个自包含的文件可以直接被PyTorch Mobile运行时加载。也就是说不需要导出ONNX不需要转换工具不需要担心某个算子被第三方引擎翻译错整个链路都在PyTorch自己的生态里。我当时唯一的顾虑是兼容性TorchScript在安卓端真的能跑目标检测模型吗搜索了一圈后发现PyTorch Mobile已经支持Android很多年了社区里有大量基于TorchScript部署YOLO的案例。而且我的自定义注意力模块本来就是用PyTorch写的TorchScript对自家算子的支持显然是最完整的。于是我决定花一个晚上试一把——结果真的就成了。4.2 导出TorchScript的完整代码与注意事项导出过程其实非常简单我用的是torch.jit.trace不是torch.jit.script。两者区别很大script会尝试把整个模型的Python控制流编译成TorchScript的IR如果你的模型里有动态分支、某些非常规写法很容易失败trace则直接拿一个示例输入跑一遍前向记录实际执行过的计算图更适合多数纯前馈的目标检测模型。import torch # 加载训练好的权重 checkpoint torch.load(best.pt, map_locationcpu) model checkpoint[model].float() model.eval() # 构造一个固定尺寸的输入 example_input torch.rand(1, 3, 640, 640) # 用 trace 导出 TorchScript traced_model torch.jit.trace(model, example_input) traced_model.save(yolov5s_custom_torchscript.pt)这段代码有四个细节必须注意一定要map_locationcpu如果模型是在GPU上保存的不在CPU上load一次就trace生成的TorchScript里可能残留GPU设备信息。一定要.float()很多训练链路会用AMP混合精度检查点里可能含有half精度的权重数据不转成float会导致CPU推理时精度异常。一定要.eval()如果不切到eval模式BatchNorm和Dropout的运行逻辑会完全不同trace出的计算图会包含训练时的行为推理结果基本是错的。一定要固定输入尺寸trace是按样例输入生成静态图的如果后面端上推理时输入尺寸和trace时不一致可能会报错。所以我统一用640×640导出。如果你的模型结构里有依赖输入值做条件判断的逻辑trace可能无法覆盖那条没被执行的路径。但这在绝大多数检测模型里不是问题YOLO这种结构就是一条笔直的前向计算链所以用trace最省事。4.3 安卓端接入Gradle依赖、加载模型、前处理后处理接下来是Android端集成。我用的PyTorch版本是1.13.1在build.gradle里加两个依赖implementation org.pytorch:pytorch_android:1.13.1 implementation org.pytorch:pytorch_android_torchvision:1.13.1加载模型的方式也很直接Module module Module.load(assetFilePath(context, yolov5s_custom_torchscript.pt));assetFilePath是一个把assets目录里的模型文件复制到应用私有目录的工具方法主要是为了得到一个绝对路径因为Module.load不接受assets路径。如果模型文件比较大加载时会有一点耗时但可以接受。接下来是预处理。我把训练阶段YOLOv5的letterbox逻辑和归一化方式在Android端重新实现了一遍把输入Bitmap按比例缩放到640×640长边保持短边用灰色像素填充填充值用114。把Bitmap像素转成float数组通道顺序为RGB。原始像素值范围是0~255需要除以255归一化到0~1。然后是推理调用Tensor inputTensor Tensor.fromBlob(floatArray, new long[]{1, 3, 640, 640}); IValue inputIValue IValue.from(inputTensor); Tensor outputTensor module.forward(inputIValue).toTensor(); float[] outputs outputTensor.getDataAsFloatArray();注意getDataAsFloatArray()返回的是一维数组需要按照模型输出的shape重新组织。我的模型输出shape是[1, 25200, 85]其中25200是YOLOv5在640×640输入下所有anchor预测框的总数85代表cx,cy,w,h,objectness,80个类别分数。这里的80指的是训练时的类别数如果你的自定义任务只有3个类别这个数字就是5类别数。后处理要在Java层自己实现先过滤掉置信度低于阈值的预测框再做NMS去掉大量重叠框最后把坐标从640×640输入图映射回原图。这一步和Python端的后处理逻辑完全一致不能偷懒用简化的替代品。4.4 为什么TorchScript能通吃我的自定义结构TorchScript能救场核心原因就是那两个字原生。我的模型里有自研注意力模块期间用到的GlobalAveragePool、Sigmoid、Mul这些操作在PyTorch的算子体系里都是基础原语PyTorch Mobile的运行时直接解释执行它们不需要经过外部转换工具重新解析图结构。换个角度说NCNN、TFLite、MNN更像是通用翻译器把PyTorch写成的内容翻译成另一门语言TorchScript则是原稿本身手机上的PyTorch Mobile就是那位原作者它当然最懂自己写的东西。整个部署链路里不再有未知算子这个概念因为所有算子都来自PyTorch自身。5. 从能跑到好用性能调优和对齐验证5.1 第一版能跑的模型慢到让人崩溃TorchScript第一次在平板上跑通的时候我很兴奋但很快就被打击了。640×640输入、float32精度、默认单线程推理跑一帧要两秒多。客户要求的是实时检测照这个速度根本没法用。虽然我的任务不需要像视频监控那样跑到30fps但至少也要5fps以上才能保证产线操作员能顺畅地扫过一个工件。那一刻我突然理解为什么很多人宁愿忍受各种转换痛苦也要上NCNN这种专门做移动端优化的框架——因为它们对CPU指令集、内存布局、线程调度的优化确实更激进。TorchScript赢在兼容性和一致性但原始性能不一定比得过专为移动端打造的推理引擎。5.2 这轮优化我到底动了哪些地方我先列了一组基线数据然后一项项做对比实验优化项输入分辨率精度模式线程数单帧耗时基线640×640float321约2200ms降低分辨率416×416float321约950ms再降分辨率320×320float321约560ms开启多线程320×320float324约180ms使用fp16320×320float164约120ms使用fp16416×416float164约210ms最终我选择在416×416的输入下运行稳定在5fps左右。虽然和NCNN理论上能达到的帧率还有差距但对于这个场景已经够用了。这个过程中最让我意外的是仅仅把输入从640降到416耗时就能少一半还多——因为计算量是输入分辨率的平方关系降幅非常明显。线程数不是越多越好。在骁龙芯片上我测试过4线程和8线程8线程并没有带来速度提升反而因为线程切换和缓存问题变得更慢。所以如果读者遇到类似问题建议自己也扫一遍线程档位别无脑拉满。fp16模式是个意外收获。PyTorch Mobile在支持fp16的CPU上可以启用半精度推理模型体积直接缩小一半速度还有提升。但要小心如果你的模型对精度极其敏感特别是微小缺陷检测fp16也可能造成漏检。我的模型测试下来影响在可接受范围内所以才敢用。5.3 量化小目标模型的教训因为TorchScript方案里面确实有量化接口我也尝试过把它转成INT8量化模型。原型验证阶段精度还可以接受一放到真实工件上就露馅很多微小划痕、浅色缺陷直接消失检测框时有时无。后来我复盘发现问题主要集中在低层特征图那里的数值范围本来就小一旦量化成8位整数很多细微纹理信息就彻底丢失了。那次之后我得出一条经验对于小目标、低对比度目标量化前一定要专门用目标类别做验证不能只看整体mAP。整体mAP可能只掉了0.02但具体到某一个难识别的小类别可能已经从0.85跌到0.4了。如果你做的是工业缺陷检测、医学影像、红外小目标这类对细节极其敏感的任务优先用fp16或直接保持fp32INT8要慎之又慎。5.4 和Python端逐条对账的验证方法性能优化做完后还有一个重要环节必须验证Android端跑出来的结果和Python端完全一致。我写了一个自动化对齐脚本把同一张测试图分别喂给Python版PyTorch模型和Android端TorchScript模型记录每个预测框的坐标、置信度和类别然后逐条对比。这一下就发现了好几个隐蔽问题我最初在Java预处理时直接用了一张非letterbox的Bitmap送入模型导致最终坐标映射时产生了偏移。后来加了letterbox但没有在NMS之后把框的坐标减掉填充区域画出来的框整体偏右下方。还有一次是Bitmap通道顺序没对齐RGB和BGR反了模型输出置信度暴跌框的位置却还大致对得上非常具有迷惑性。最靠谱的对账做法是准备20张典型图片包含不同光照条件、不同目标大小、不同缺陷类型在Python端导出每个框的完整信息存成JSON在Android端也输出同样的JSON然后用脚本比较它们的IoU和置信度差异。我最后设定了两个阈值框中心偏移不超过2个像素、置信度相差不超过0.05。全部满足才算通过。6. 部署完成后的框架选型心得6.1 一张表总结四个方案的实战感受整个流程走完我把四个方案的实际体验整理成一张表方案转换链路复杂度自定义结构兼容性端侧调优难度调试一致性最终结局NCNN高依赖ONNX中间层低遇自定义算子很容易卡住中工具链丰富一般败在算子兼容TFLite很高多层工具链组合低到中高量化参数难调一般败在转换和精度损失MNN中较高但转换工具版本影响大中较差各后端结果不一致败在一致性TorchScript极低PyTorch原生高基本无缝偏低官方接入成熟高与Python端几乎一致最终采用这条对比给了我一个很重要的启发当你手里已经有一个PyTorch模型时先别急着去做各种格式转换先利用PyTorch自己的移动端生态把整条链路跑通再考虑性能优化。跑通后再迁移到NCNN这类高性能引擎你的目标会更明确——你知道哪些算子必须解决、哪些精度差异可以接受、性能瓶颈到底在哪。6.2 如果再让我部署一次目标检测模型我会怎么走如果现在让我重新部署一个目标检测模型到安卓端我的路线会完全不一样第一步先看模型结构再选框架。如果模型里有比较新的算子、自定义模块、注意力机制直接默认选TorchScript/PyTorch Mobile不要浪费时间在第三方转换工具链上。第二步先跑通Demo再优化性能。哪怕初始帧率只有个位数也先把端到端流程走完确立一条可用的基准线。没有基准线的一切性能优化都是空中楼阁。第三步性能不够再考虑迁移到NCNN/TFLite。迁移前先确认你的模型结构在目标框架的算子支持范围内或者做好为未知算子写自定义层的心理准备。第四步量化排最后。先用fp32跑通再用fp16验证精度最后才根据需求评估要不要INT8。小目标场景尽量别碰INT8。另外提一句PyTorch Mobile现在也在演进官方主力正在往ExecuTorch方向迁移。ExecuTorch面对的问题比TorchScript更复杂会更贴合移动端和边缘设备。但TorchScript这套方案对于存量项目来说仍然非常稳定、可靠至少在我那个车间平板上的项目到现在还在用着。这台平板如今还在产线上服役每天跟着工人扫过一批又一批零件。每次路过它看到屏幕上实时框出缺陷位置的时候我都会想起那段对着红色报错日志发呆的日子。现在再有人问我安卓端部署目标检测模型难不难我会说如果选对路不难如果一开始就跳进转换工具链的坑里轻则熬夜重则怀疑人生。
返回列表