乐趣区

关于人工智能:BaseDet-走过开发的弯路

jiezi

作者：王枫 | 旷视算法研究员 

收到 MegEngine 团队的邀请来写这篇稿子，本意是想让我介绍一下 BaseDet（一个基于 MegEngine 写成的指标检测仓库，相似 detectron2 之于 pytorch）。因为大部分介绍框架的稿件总是在抓着一些代码中的 feature 疯狂介绍，而我自己并不是很喜爱这种格调（因为这些内容很像是把文档翻译成了文章），所以本文在介绍 BaseDet 之外，分享在实现 BaseDet 过程中面临的问题和思考。这些内容波及的范畴比拟广，有对于深度学习框架、软件工程和开源我的项目等诸多内容；而这些问题和思考当然也不仅仅来源于 BaseDet，同时也蕴含 MegEngine 团队在不断完善各种性能时候的踩坑与反思。

本文不会介绍具体的检测模型是怎么样的，也不会介绍实现时候的应用的提点 trick 或者具体的细节，如果你对细节感兴趣，能够参考一个我之前写的炼丹细节 blog。然而如果你关怀“当初的各类训练框架是怎么设计的”，“为什么会有这种设计”之类的问题，本文或者能够帮你了解一些外在的起因。

mmdet 与 detectron2

提到检测框架，简直任何一个做过检测相干钻研的的人都用过 mmdet 或者 detectron2 其中的一个，而做检测利用相干的人则十分偏向于应用 YOLO 系列的各个框架。在文章前面咱们会聊，钻研和利用选用不同的办法的景象是存在一些比拟粗浅的起因的。

在那之前，咱们还是聊回 BaseDet 和 mmdet/detectron2 这两大框架的一些分割。BaseDet 其实借鉴了一些 mmdet 和 detectron2 中精华的设计和理念：Trainer 和 hook。

Trainer 定义了训练逻辑的最外围组件：模型、优化器和 dataloader，简直大部分的训练场景都能够用着三个组件实现，也就是上面的逻辑：

data = next(dataloader)
loss = model(data)
loss.backward()
solver.step()

在 mmdet/detectron2/BaseDet 外面，所有的训练外围流程都是下面这个十分简短的函数，而至于 dataloader，model 和 solver 这三个常常发生变化的对象，通常是借助工厂模式的 build 办法产生的，要改哪个局部，用户只须要本人 build 就行了。

Hook 则是训练逻辑的内涵，因为在训练过程中经常会插入一些特定的需要，比方训练的一些数据 log 进 tensorboard/wandb、每训练完几个 epoch 就对模型进行一下测试、保留训练的断点等一些性能，这些性能以及对应的延长性能都依赖于 hook 的引入。

了解了 Trainer 和 Hook 的概念之后，用户其实就能够很容易对本人的需要做裁减，而诸如 dataloader、model、solver 都是能够本人 build 进去的，为了用户可能把 mmdet/detectron2 当作一个仓库应用，这两个框架都提供了注册机制（registry）。

须要留神的是，hook 和 registry 的引入都是基于这样的 trade-off：就义掉一部分用户的应用门槛，换取框架的灵活性的晋升，把一部分对于保护人员的艰难转移给了一部分用户。对于 YOLO 系列的框架（比方 YOLOv5/YOLOX 等）就不会存在这样的 trade-off：一方面模型很少，另一方面就是大部分用户还是偏向于 clone 下来本人魔改 code，对于这样的用户群体来说，晓得在哪里批改就肯定能产生成果是最重要的，此时 KISS 准则（Keep It Simple and Stupid）就显得分外重要。

MegEngine 和 DTR

BaseDet 是基于 MegEngine 的一个检测框架，如果要聊 feature，实质上也是聊 MegEngine 的 feature，毕竟 BaseDet 只是帮忙用户实现一些根本的训练任务，乏味的 feature 还是由底层框架反对的，所以这个局部咱们来聊一聊 MegEngine。

为了用户的迁移性，MegEngine 在一些 API 上和 numpy 做了对齐，这点上和 google 的 jax 是比拟相似的，益处是因为 numpy 的 api 比较稳定且 well-known；而 MegEngine 在 module 的上的设计比拟靠近 torch，因为用户对于 torch 的 module 的用法是绝对相熟的。对于大部分 torch 用户，要转 MegEngine 还是绝对比拟丝滑的，最须要留神的点就是：在 MegEngine 外面，autograd 是由一个叫做 GradManager 的 class 管制的，有点相似 tensorflow 的 GradientTape，这样做的益处在于不便管制资源的治理，不容易像 torch 一样呈现奇怪的内存透露景象（对于这个景象感兴趣的同学，能够参考之前我写的另一个 blog）。

我集体最喜爱的 MegEngine 的 feature 是由 @圆角骑士魔理沙提出来的 DTR（Dynamic Tensor Rematerialization，举荐去看原文），以 FCOS 的 baseline 为例，在 2080Ti 上单卡训练，不开 DTR batchsize 只能开到 8，关上 DTR 的状况下，batchsize 能翻一倍开到 16（当然训练速度也会变慢）。

当然，有很多实现细节是原文没有思考的，依据 engine 团队的整顿，也在这里分享一些坑点（倡议看完论文再来看这里的坑点，了解更粗浅一些）：

  1. 多卡反对。原始论文没思考这个问题，其实说起来解决办法很简略，就是无脑把须要做 send/recv 通信的 tensor 当成 immutable 的，不要 drop 就好了。
  2. 显存碎片常常会导致算法不实用，实际上估值函数须要与内存分配器联动。这里咱们为了不便了解举个例子。假如显存的状态是有 200Mb 能够自在应用，其排布形式是[A(90M) B(10M) C(90M) D(10M)]，其中 A、B、C、D 都是 tensor，括号外面是 tensor 须要的显存大小。假如有新的 tensor E 须要 15M 的空间，假如 DTR 默认算进去是 drop 掉 B 和 D，然而因为显存不间断，此时还须要 drop 掉 A 或者 C，那么一开始 drop 掉 B 和 D 的行为就很不划算，不如一开始就 drop 掉 A 或者 C，所以说估值函数实际上是须要和内存分配器做联动。在 pytorch 外面很难获取到当初各个 blob 的申请状况，而 mge 里的显存分配器设计的比拟洁净，申请开释也都有对立的中央，所以 DTR 这个机制实现的也绝对洁净一些。
  3. 波及跨 iter 操作的时候会有一些麻烦，比方 ema 中须要进行非凡解决，能够参考 BaseDet 外面的示例 code。呈现问题的起因在于：诸如 ema 这样的操作，通常会使得 tensor 的计算历史成为一个有限长（和训练长度一样）的货色，而 DTR 就会把历史上用到的 tensor 都记下来（重算过程须要应用），这就会导致呈现透露景象。
  4. 原始论文里收到 pytorch 限须要手动填阈值，大部分用户并不是很喜爱这种调用形式，最初在 MegEngine 外面应用的是一个自适应的阈值。对于用户来说，只须要在 code 外面加上 mge.dtr() 就能简略开启性能了。

不同用户的不同需要

在旷视外部有一个很棒的帖子，讲的是用户通常只会用到软件中 15% 的性能，而不同类型的用户应用的往往是同一个软件中那不同的 15% 局部。在实现 BaseDet 的过程中，我接触到了不同的用户人群，理解到这些人群对于框架的不同需要。举个例子：

  • 钻研人员：灵便，但同时有须要的性能的时候能够简略关上（比方 ema）。喜爱 pytorch-lightning/timm 这种 lite 的货色，关怀训练 / 评测逻辑，训练进去的模型点数越高越好。
  • 产品研发：关怀的重要的参数 可能简略配置，不便交付。喜爱 onnx/torchscript 这种两头产物，不关怀训练评测模型的逻辑，像保姆一样帮他们搞个 demo 走通流程最好。
  • 深度学习框架研发：须要简略就能跑起来的仓库，不便追溯问题。下层爱咋写咋写，爱咋封装咋封装，喜爱训练框架提供诸如保留 crashing context、profiler、benchmark 等性能。

所以诸如 mmdet/detectron2 这类框架都是反对简略的 yaml config 和 lazy eval 的性能的，看起来可能有些矛盾，然而这种做法可能满足不同群体的需要。

后面提到过，做检测利用相干的人则十分偏向于应用 YOLO 系列的各个框架，一部分起因就是大部分用户是间接 clone 下来仓库间接改 code 的，所以在这些框架中，很少提供诸如 registry 和 hook 这类概念，因为这些概念自身并没有提供灵活性，反而引入了多余的概念。

因为 BaseDet 自身是为了辅助产品而存在的，所以是基于 product first 的准则而设计开发的，也就不可避免地在应用体验上存在一些 bias，开源进去的目标其实就是为了纠正这种 bias，还能给 MegEngine 的用户提供一种 code 参考，心愿社区可能给予一些适当的反馈，这些反馈也是 codebase 后退的方向。

BaseDet 应用示例：https://studio.brainpp.com/project/28826?name=BaseDet%E4%BD%BF%E7%94%A8%E7%A4%BA%E4%BE%8B

后记

留下来一段话，送给这世界上违心破费工夫精力去 maintain 我的项目的开发人员，也是我这一段时间来的粗浅感悟：任何一段 code 都值得一直破费工夫去打磨，然而打磨之后的 code 并不是真正的产出，关键在于过程中的思考和学习。不应该和本人保护的的仓库适度绑定，总有一些更重要的事件在等着你。

更多 MegEngine 信息获取，您能够：查看文档、和 GitHub 我的项目，或退出 MegEngine 用户交换 QQ 群：1029741705。欢送参加 MegEngine 社区奉献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停。

退出移动版