BaseDet: 走过开发的弯路

[MegEngineBot]

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

收到 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，荣誉证书、定制礼品享不停。