这两年，端侧AI发展得很快。

以前，AI模型都跑在云端，手机和电脑只是个客户端，负责把数据传到云端，然后接收结果。但现在，越来越多的AI模型开始跑在端侧，也就是手机、电脑、IoT设备这些终端设备上。

端侧AI有很多好处：延迟低，不需要网络传输；隐私好，数据不需要传到云端；成本低，不需要昂贵的云端服务器。所以，端侧AI正在快速普及，越来越多的应用开始集成端侧AI功能。

但端侧AI的开发，和普通的应用开发很不一样。它涉及到模型部署、性能优化、内存管理、硬件加速等很多复杂的技术。很多端侧AI项目，因为赶进度、缺经验，代码写得很烂，性能差、内存高、难维护。

我所在的团队，这两年做了好几个端侧AI项目。其中有一个项目，因为早期代码写得太烂，后期维护和优化都非常困难。最后，我们下决心做了一次彻底的代码重构，把烂代码改成了优雅、高效、可维护的代码。

这篇文章，我想分享一下这次端侧AI项目代码重构的实战经验。从性能优化、内存管理到模型部署、代码架构，聊聊如何把端侧AI项目的烂代码改造成优雅、高效、可维护的代码。

如果你也在做端侧AI开发，或者打算做，希望这篇文章能给你一些参考，让你少踩一些坑。

先说明一下，我们的项目主要是在移动端（Android和iOS）上部署端侧AI模型，用的推理框架是ONNX Runtime和Core ML。不同的平台和框架，具体实现会有差异，但重构的思路和方法是相通的。

为什么要重构

在讲重构之前，先说说我们为什么要重构。

那个端侧AI项目，是我们团队的第一个端侧AI项目。刚开始的时候，大家都没什么经验，加上项目时间紧，就怎么快怎么来，先把功能跑通再说。

结果，项目上线之后，问题就来了：

第一，性能差。模型推理速度慢，用户体验不好。特别是在中低端手机上，推理一次要好几秒，用户等得不耐烦。

第二，内存高。应用占用的内存很大，经常导致内存不足，在一些内存小的手机上甚至会崩溃。

第三，代码乱。代码结构混乱，耦合严重，一个文件几千行，一个函数几百行。改一个地方，可能会影响很多地方，bug不断。

第四，难维护。因为代码乱，新人接手很困难，老员工改起来也头疼。每次加新功能或者修bug，都要花很长时间理解代码。

第五，难扩展。因为耦合严重，想加一个新模型或者新功能，都要改很多地方，很容易引入新的bug。

这些问题，随着项目的迭代，越来越严重。最后，团队达成共识：必须重构，否则这个项目就没法维护了。

当然，重构不是一拍脑袋就决定的。我们做了充分的评估：重构的成本有多大，风险有多高，收益有多少，时间怎么安排。评估之后，我们觉得收益大于成本和风险，于是就开始了重构。

重构前的准备

和任何重构一样，端侧AI项目的重构，准备工作也非常重要。

准备一：全面了解现有代码

重构之前，首先要全面了解现有代码。

我们花了一周时间，把整个项目的代码都过了一遍，搞清楚了：

  • 代码的整体结构和模块划分
  • 模型加载和推理的流程
  • 性能瓶颈在哪里
  • 内存占用高的原因是什么
  • 代码中最乱、耦合最严重的部分是哪些
  • 有哪些技术债务需要偿还

了解了这些，我们才能有针对性地制定重构计划，知道哪些地方是重点，哪些地方可以先放一放。

准备二：建立性能基准

重构的目标之一是提升性能，所以重构之前，必须建立性能基准。

我们用 profiling 工具，对现有代码做了全面的性能分析，记录了：

  • 模型加载时间
  • 单次推理时间
  • 内存峰值占用
  • CPU和GPU使用率
  • 不同机型上的性能表现

有了这些基准数据，重构之后才能对比，知道性能有没有提升，提升了多少。

同时，我们也建立了自动化的性能测试，每次重构之后都跑一遍，确保性能不会下降。

准备三：确保有测试

重构的前提是有测试。如果没有测试，重构就是在裸奔。

端侧AI项目的测试，和普通项目不太一样。除了常规的单元测试和集成测试，还需要：

  • 模型输出一致性测试：确保重构之后，模型的输出和重构之前一致
  • 性能测试：确保重构之后，性能没有下降
  • 兼容性测试：确保在不同机型、不同系统版本上都能正常运行

我们在重构之前，补了很多测试，特别是模型输出一致性测试。这个测试非常重要，它能确保我们在重构的过程中，不会不小心改变模型的行为。

准备四：制定重构计划

准备工作做好之后，我们制定了详细的重构计划。

重构计划包括：

  • 重构的目标和范围
  • 重构的步骤和顺序
  • 每个步骤的时间安排
  • 每个步骤的验收标准
  • 风险评估和应对措施

我们没有选择一次性大重构，而是选择了增量式重构。先重构最核心、问题最严重的部分，然后逐步扩展到其他部分。每完成一个小的重构，就测试、验证，确保没有问题，再进行下一步。

这样，重构的风险就小很多，即使出了问题，也容易定位和回滚。

重构第一步：模型管理模块

我们重构的第一步，是模型管理模块。

模型管理是端侧AI项目最核心的部分，也是原来代码最乱的部分。原来的代码里，模型的加载、卸载、缓存、版本管理，都混在一起，耦合严重，很难维护。

我们把模型管理模块拆分成了几个独立的部分：

第一，模型加载器。负责模型的加载，包括从本地文件加载、从资源加载、从网络下载加载。加载器有统一的接口，不同来源的模型用不同的加载器实现。

第二，模型缓存。负责模型的缓存管理，包括内存缓存和磁盘缓存。缓存有统一的策略，比如LRU淘汰、最大缓存大小限制等。

第三，模型版本管理。负责模型的版本管理，包括版本检查、版本升级、版本回滚。确保用户总是使用最新的模型，同时支持回滚到旧版本。

第四，模型生命周期管理。负责模型的生命周期，包括模型的创建、初始化、使用、销毁。确保模型在不需要的时候能及时释放资源，避免内存泄漏。

拆分之后，每个部分都有明确的职责，代码清晰了很多，也更容易测试和维护。

同时，我们也优化了模型加载的性能。原来的代码，每次使用模型都要重新加载，很慢。重构之后，我们加了模型缓存，常用的模型加载一次之后就缓存在内存里，下次使用直接从缓存取，大大提升了性能。

重构第二步：推理引擎模块

重构的第二步，是推理引擎模块。

推理引擎是端侧AI项目性能最关键的部分。原来的代码里，推理逻辑写得很随意，没有统一的接口，不同模型用不同的推理方式，性能也很差。

我们把推理引擎模块重构为统一的架构：

第一，统一的推理接口。定义了一个统一的推理接口，所有模型都通过这个接口进行推理。接口包括模型加载、输入预处理、推理执行、输出后处理等步骤。这样，不管是什么模型，用起来都是一样的，大大降低了使用成本。

第二，推理后端抽象。不同的平台和硬件，可能用不同的推理后端，比如Android上用ONNX Runtime或者NNAPI，iOS上用Core ML。我们把推理后端抽象出来，定义统一的后端接口，不同的后端用不同的实现。这样，上层代码不需要关心具体用的是什么后端，后端可以根据平台和硬件自动选择。

第三，硬件加速。端侧AI的性能，很大程度上取决于能不能利用硬件加速。我们在重构的时候，充分利用了各个平台的硬件加速能力，比如Android上的NNAPI、GPU加速，iOS上的Core ML、Metal。硬件加速能让推理速度提升好几倍。

第四，推理优化。除了硬件加速，我们还做了很多推理优化，比如模型量化、算子融合、内存复用等。这些优化能进一步提升推理速度，降低内存占用。

第五，异步推理。原来的代码里，推理是同步的，会阻塞UI线程，导致界面卡顿。重构之后，我们把推理改成了异步的，在后台线程执行，通过回调返回结果。这样，界面就不会卡顿了，用户体验好了很多。

重构之后，推理引擎的性能提升了很多。在中低端手机上，推理时间从原来的好几秒降到了一秒以内，内存占用也降低了一半以上。

重构第三步：数据预处理和后处理模块

重构的第三步，是数据预处理和后处理模块。

端侧AI模型，对输入数据的格式有严格的要求。比如，图像模型需要把图片resize到固定尺寸，归一化到特定的范围，转换为特定的颜色空间。输出也需要后处理，比如检测模型需要解析输出的边界框和类别，分类模型需要计算概率和排序。

原来的代码里，预处理和后处理逻辑散落在各个地方，每个模型都有自己的一套，重复代码很多，而且经常写错。

我们把预处理和后处理模块重构为统一的、可复用的架构：

第一，预处理管道。定义了一个预处理管道，可以把多个预处理步骤组合起来，比如resize、crop、normalize、color convert等。每个步骤都是独立的、可复用的，可以根据不同模型的需求，灵活组合。

第二，后处理管道。同样，定义了后处理管道，把多个后处理步骤组合起来，比如decode、nms、sort、filter等。不同模型可以根据自己的输出格式，灵活组合后处理步骤。

第三，数据格式转换。端侧AI涉及到很多数据格式的转换，比如图片格式转换（Bitmap到ByteBuffer）、张量格式转换（NCHW到NHWC）、数据类型转换（float到uint8）等。我们把这些转换封装成统一的工具函数，避免重复造轮子。

第四，性能优化。预处理和后处理，有时候也会成为性能瓶颈，特别是图像处理相关的操作。我们对这些操作做了优化，比如利用多线程、SIMD指令、硬件加速等，提升预处理和后处理的速度。

重构之后，预处理和后处理的代码清晰了很多，重复代码大大减少，也更容易维护和扩展。加新模型的时候，只需要组合已有的预处理和后处理步骤，不需要再从头写一遍。

重构第四步：内存管理

重构的第四步，是内存管理。

端侧AI项目，内存是一个大问题。AI模型本身就很大，加载到内存里要占很多空间；推理过程中，中间的张量也要占很多内存。如果内存管理不好，很容易导致内存不足，甚至崩溃。

原来的代码里，内存管理很随意，模型加载了不卸载，张量创建了不释放，内存泄漏很严重。

我们在重构的时候，对内存管理做了全面的优化：

第一，模型生命周期管理。确保模型在不需要的时候能及时卸载，释放内存。我们用引用计数的方式来管理模型的生命周期，当没有地方使用模型的时候，自动卸载。

第二，张量内存复用。推理过程中，会创建很多中间张量。原来的代码里，每个张量都单独分配内存，浪费很大。重构之后，我们实现了张量内存池，不同的张量可以复用同一块内存，大大降低了内存占用。

第三，图片内存管理。图像处理是内存大户。我们对图片的内存做了专门的管理，比如及时回收不需要的图片，复用图片缓冲区，避免频繁的内存分配和释放。

第四，内存监控和预警。我们加了内存监控功能，实时监控应用的内存使用情况。当内存使用超过阈值的时候，会自动卸载一些不常用的模型，释放内存，避免崩溃。

第五，内存泄漏检测。我们用内存泄漏检测工具，定期检查代码中的内存泄漏，及时发现和修复。

重构之后，应用的内存占用降低了很多，内存泄漏的问题也基本解决了。在内存小的手机上，也能稳定运行，不会轻易崩溃。

重构第五步：代码架构和模块化

重构的第五步，是代码架构和模块化。

前面几步，我们重构了各个核心模块。这一步，我们对整个项目的代码架构做了梳理和优化，让整个项目的结构更清晰，耦合更低，更容易维护和扩展。

我们做了以下几件事情：

第一，分层架构。我们把项目分成了几个清晰的层次：UI层、业务逻辑层、AI能力层、基础设施层。每一层都有明确的职责，层与层之间通过接口通信，依赖关系清晰。UI层不直接操作AI模型，而是通过业务逻辑层调用；业务逻辑层不关心AI的具体实现，而是通过AI能力层的接口调用。

第二，模块化。我们把项目拆分成了多个独立的模块，每个模块负责一个特定的功能。比如，模型管理模块、推理引擎模块、图像处理模块、音频处理模块等。每个模块都是独立的，可以单独编译、单独测试、单独复用。模块之间通过接口通信，耦合很低。

第三，依赖注入。我们用依赖注入的方式，来管理模块之间的依赖关系。这样，模块之间的耦合更低，更容易测试，也更容易替换实现。比如，推理引擎的实现可以替换，上层代码不需要改。

第四，接口和实现分离。每个模块都定义了清晰的接口，接口和实现分离。上层代码只依赖接口，不依赖具体实现。这样，实现可以灵活替换，比如不同平台用不同的实现，测试的时候可以用mock实现。

第五，代码规范。我们制定了统一的代码规范，包括命名规范、注释规范、格式规范、设计原则等。所有代码都遵循统一的规范，看起来整齐划一，可读性大大提高。

重构之后，整个项目的代码结构清晰了很多，耦合大大降低。新人接手项目，能很快理解代码结构；加新功能或者修bug，也不需要改很多地方；代码的可维护性和可扩展性都大大提升。

重构中的性能优化技巧

在重构的过程中，我们积累了一些端侧AI性能优化的技巧，这里分享给大家。

第一，模型量化。模型量化是把模型的权重从float32转换成int8或者float16，能大大减小模型大小，提升推理速度，降低内存占用。量化对精度的影响通常很小，但性能提升很大。端侧AI项目，几乎都应该做模型量化。

第二，算子融合。算子融合是把多个小算子合并成一个大算子，减少中间结果的内存读写和kernel launch开销。比如，把conv+bn+relu融合成一个算子。算子融合能显著提升推理速度，降低内存占用。

第三，内存复用。推理过程中，中间张量的内存可以复用。比如，前一层的输出张量，在计算完下一层之后就不需要了，它的内存可以被后面的张量复用。内存复用能大大降低推理过程中的内存峰值。

第四，多线程。端侧设备通常有多个CPU核心，可以利用多线程来加速推理。比如，把模型的不同部分分配到不同的线程上并行计算，或者用多线程做预处理和后处理。多线程能充分利用CPU资源，提升推理速度。

第五，硬件加速。尽可能利用硬件加速能力，比如GPU、NPU、DSP等。不同的平台有不同的硬件加速接口，比如Android的NNAPI，iOS的Core ML和Metal。硬件加速能让推理速度提升好几倍，是端侧AI性能优化最重要的手段之一。

第六，预处理和后处理优化。预处理和后处理有时候也会成为性能瓶颈，特别是图像处理相关的操作。可以用多线程、SIMD指令、硬件加速等方式来优化预处理和后处理的速度。

第七，模型裁剪和蒸馏。如果模型太大或者太慢，可以考虑模型裁剪（去掉不重要的参数）和知识蒸馏（用大模型训练小模型）。这些方法能在精度损失很小的情况下，大大减小模型大小，提升推理速度。

这些技巧，我们在重构中都用到了，效果都很不错。当然，具体用哪些技巧，要根据项目的实际情况来决定。

重构中的坑和经验

在重构的过程中，我们也踩了很多坑，积累了一些经验。

坑一：模型输出不一致

重构的时候，最担心的就是模型输出不一致。因为推理逻辑改了，预处理和后处理也改了，很容易导致模型的输出和原来不一样。

我们遇到过这个问题。重构之后，模型的输出和原来有细微的差别。虽然差别很小，但在某些场景下会导致结果不对。

后来，我们建立了严格的模型输出一致性测试。重构之前，我们用原来的代码跑了大量的测试用例，记录了模型的输出。重构之后，用同样的测试用例跑，对比输出是否一致。如果不一致，就排查原因，直到完全一致。

有了这个测试，我们就放心多了。每次重构之后，跑一遍测试，就能知道模型输出有没有问题。

坑二：不同机型的兼容性问题

端侧AI的一个大问题，是不同机型的兼容性。不同的手机，芯片不同，系统版本不同，硬件加速能力不同，同样的代码，在不同手机上的表现可能完全不一样。

我们在重构的时候，遇到过很多兼容性问题。比如，某个硬件加速在某些机型上有bug，导致推理结果错误；某个优化在某些机型上反而更慢；某个API在旧版本系统上不支持。

后来，我们建立了完善的兼容性测试。我们准备了不同品牌、不同价位、不同系统版本的测试机，每次重构之后都在所有测试机上跑一遍，确保兼容性。

同时，我们也做了降级机制。如果硬件加速在某个机型上有问题，就自动降级到CPU推理；如果某个优化在某个机型上更慢，就自动关闭这个优化。这样，即使有兼容性问题，也能保证应用正常运行，只是性能差一些。

坑三：重构期间的功能开发

重构期间，项目不能停，新功能还要继续开发。这就带来了一个问题：重构和功能开发并行，代码冲突很多，很容易出问题。

我们的做法是，把重构分成很多小的步骤，每个步骤都是独立的、可合并的。功能开发在主分支上进行，重构在单独的分支上进行，每个小的重构步骤完成之后，就合并到主分支。这样，代码冲突就少很多，也不容易出问题。

同时，我们也和团队成员充分沟通，让大家知道重构的计划和进度，在写新代码的时候遵循新的架构和规范，避免产生新的技术债务。

坑四：过度设计

重构的时候，很容易过度设计。为了追求"优雅"，加了很多不必要的抽象和层，结果代码变得更复杂，更难维护。

我们在重构的时候，也犯过这个错误。比如，为了"可扩展性"，加了很多接口和抽象类，但实际上根本不需要这么多扩展点。结果代码结构变得很复杂，新人看了都晕。

后来，我们意识到了这个问题，开始做减法。遵循KISS原则（Keep It Simple, Stupid），只在真正需要的地方做抽象和扩展。能简单实现的，就简单实现，不要为了可能的未来需求而过度设计。

记住，重构的目的是让代码更简单、更清晰、更容易维护，而不是更复杂、更"高级"。

重构的效果

经过几个月的努力，我们终于完成了这次重构。重构的效果，超出了我们的预期。

第一，性能大幅提升。模型加载时间降低了60%，推理时间降低了50%以上，应用启动速度也快了很多。在中低端手机上，用户体验好了很多。

第二，内存大幅降低。应用的内存峰值占用降低了40%，内存泄漏的问题基本解决。在内存小的手机上，也能稳定运行，不会轻易崩溃。

第三，代码质量大幅提升。代码结构清晰，耦合低，规范统一。代码行数减少了30%，但可读性和可维护性大大提升。

第四，开发效率提升。因为代码清晰了，耦合低了，加新功能和修bug的速度快了很多。新人接手项目，也能很快上手。

第五，团队士气提升。重构之前，大家都很痛苦，改代码改到崩溃。重构之后，代码清晰了，大家写代码也更有信心了，团队士气提升了很多。

当然，重构也不是没有成本。我们花了几个月的时间，投入了好几个人力，期间也遇到了很多问题和挑战。但总的来说，收益远远大于成本。这次重构，让这个项目"活"了过来，为后续的发展打下了坚实的基础。

给端侧AI开发者的建议

最后，给正在做或者打算做端侧AI开发的朋友一些建议。

第一，从一开始就重视代码质量。不要因为赶进度就写烂代码，技术债务迟早要还的，而且越晚还，代价越大。从项目一开始，就建立好的代码架构和规范，重视代码质量，能省很多后期的麻烦。

第二，性能和内存要从一开始就考虑。端侧AI，性能和内存是生命线。不要等功能都做完了才去优化性能和内存，那样成本很高。从一开始，就把性能和内存作为重要的设计目标，在架构设计和代码实现中就考虑进去。

第三，充分利用硬件加速。端侧AI的性能，很大程度上取决于硬件加速。从一开始，就设计好硬件加速的架构，充分利用各个平台的硬件加速能力。不要等性能不够了才去加硬件加速，那样重构的成本很高。

第四，建立完善的测试。端侧AI的测试很重要，特别是模型输出一致性测试、性能测试、兼容性测试。从一开始就建立完善的测试，能帮你及时发现问题，避免问题积累。

第五，持续重构。重构不是一次性的事情，而是持续的过程。在日常开发中，就要不断地优化代码，偿还技术债务。不要等代码烂到没法维护了才去重构，那样代价太大。小步快跑，持续重构，是最好的方式。

这些建议，是我们用血泪换来的经验，希望能帮到大家。

写在最后

端侧AI正在快速普及，越来越多的应用开始集成端侧AI功能。但端侧AI的开发，和普通的应用开发很不一样，它涉及到很多复杂的技术，也很容易产生技术债务。

代码重构，是解决技术债务、提升代码质量的有效方式。但重构不是一件容易的事情，它需要充分的准备、清晰的计划、正确的方法，也需要团队的共识和耐心。

我们这次端侧AI项目的重构，虽然过程很辛苦，踩了很多坑，但最终的结果是好的。重构之后，项目的性能、内存、代码质量、开发效率都有了大幅提升。

如果你也在做端侧AI开发，也面临着代码烂、性能差、难维护的问题，我建议你认真考虑一下重构。虽然重构有成本，但长期来看，收益远远大于成本。

当然，最好的方式，还是从一开始就重视代码质量，避免产生太多技术债务。但如果已经产生了，也不要怕，勇敢地去重构，把烂代码改成优雅代码。

最后，用一句话来结束这篇文章："好的代码不是写出来的，而是改出来的。端侧AI的代码更是如此，需要不断地优化、重构，才能做到既高效又优雅。"

愿你在端侧AI的开发道路上，写出高效、优雅、可维护的代码。