博客

BaseDet: 走过开发的弯路
作者:王枫 发布日期:2022/12/15

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

查看更多

上一篇: 分享

Awesome MegEngineer 英雄招募帖,开源社区专属权益等你来领

下一篇: 分享

如何写一个深度学习编译器

相关推荐
  • 金枪鱼之夜:MegEngine 框架设计

    2020/08/23

  • 旷视天元Beta版核心技术升级全面解读

    2020/07/24

  • MegEngine Windows Python wheel 包减肥之路

    2022/01/05

  • MegEngine 架构系列:静态内存分析

    2021/10/21

  • MegCC 新版本来啦!新增 Benchmark 等工具,性能大幅提升!有奖征文活动同步启动!

    2023/01/12