作者：王楓 | 曠視算法研究員?@Fatescript

收到 MegEngine 團(tuán)隊(duì)的邀請來寫這篇稿子，本意是想讓我介紹一下?BaseDet（一個(gè)基于?MegEngine?寫成的目標(biāo)檢測倉庫，類似 detectron2 之于 pytorch）。因?yàn)榇蟛糠纸榻B框架的稿件總是在抓著一些代碼中的 feature 瘋狂介紹，而我本人并不是很喜歡這種風(fēng)格（因?yàn)檫@些內(nèi)容很像是把文檔翻譯成了文章），所以本文在介紹 BaseDet 之外，分享在完成 BaseDet 過程中面臨的問題和思考。這些內(nèi)容涉及的范圍比較廣，有關(guān)于深度學(xué)習(xí)框架、軟件工程和開源項(xiàng)目等諸多內(nèi)容；而這些問題和思考當(dāng)然也不僅僅來源于 BaseDet，同時(shí)也包含 MegEngine 團(tuán)隊(duì)在不斷完善各種功能時(shí)候的踩坑與反思。

本文不會介紹具體的檢測模型是怎么樣的，也不會介紹實(shí)現(xiàn)時(shí)候的使用的提點(diǎn) trick 或者具體的細(xì)節(jié)，如果你對細(xì)節(jié)感興趣，可以參考一個(gè)我之前寫的煉丹細(xì)節(jié)?blog。但是如果你關(guān)心“現(xiàn)在的各類訓(xùn)練框架是怎樣設(shè)計(jì)的”，“為什么會有這種設(shè)計(jì)”之類的問題，本文或許可以幫你理解一些內(nèi)在的原因。

mmdet 與 detectron2

提到檢測框架，幾乎任何一個(gè)做過檢測相關(guān)研究的的人都用過 mmdet 或者 detectron2 其中的一個(gè)，而做檢測應(yīng)用相關(guān)的人則非常傾向于使用 YOLO 系列的各個(gè)框架。在文章后面我們會聊，研究和應(yīng)用選用不同的方法的現(xiàn)象是存在一些比較深刻的原因的。

在那之前，我們還是聊回 BaseDet 和 mmdet/detectron2 這兩大框架的一些聯(lián)系。BaseDet 其實(shí)借鑒了一些 mmdet 和 detectron2 中精髓的設(shè)計(jì)和理念：Trainer 和 hook。

Trainer 定義了訓(xùn)練邏輯的最核心組件：模型、優(yōu)化器和 dataloader，幾乎大部分的訓(xùn)練場景都可以用著三個(gè)組件完成，也就是下面的邏輯：

在 mmdet/detectron2/BaseDet 里面，所有的訓(xùn)練核心流程都是上面這個(gè)非常簡短的函數(shù)，而至于 dataloader，model 和 solver 這三個(gè)經(jīng)常發(fā)生變化的對象，通常是借助工廠模式的 build 方法產(chǎn)生的，要改哪個(gè)部分，用戶只需要自己 build 就行了。

Hook 則是訓(xùn)練邏輯的外延，因?yàn)樵谟?xùn)練過程中常常會插入一些特定的需求，比如訓(xùn)練的一些數(shù)據(jù) log 進(jìn) tensorboard/wandb、每訓(xùn)練完幾個(gè) epoch 就對模型進(jìn)行一下測試、保存訓(xùn)練的斷點(diǎn)等一些功能，這些功能以及對應(yīng)的延伸功能都依賴于 hook 的引入。

理解了 Trainer 和 Hook 的概念之后，用戶其實(shí)就可以很容易對自己的需求做擴(kuò)充，而諸如 dataloader、model、solver 都是可以自己 build 出來的，為了用戶能夠把 mmdet/detectron2 當(dāng)作一個(gè)倉庫使用，這兩個(gè)框架都提供了注冊機(jī)制（registry）。

需要注意的是，hook 和 registry 的引入都是基于這樣的 trade-off：犧牲掉一部分用戶的使用門檻，換取框架的靈活性的提升，把一部分對于維護(hù)人員的困難轉(zhuǎn)移給了一部分用戶。對于 YOLO 系列的框架（比如 YOLOv5/YOLOX 等）就不會存在這樣的 trade-off：一方面模型很少，另一方面就是大部分用戶還是傾向于 clone 下來自己魔改 code，對于這樣的用戶群體來說，知道在哪里修改就一定能產(chǎn)生效果是最重要的，此時(shí) KISS 原則（ Keep It Simple and Stupid ）就顯得格外重要。

MegEngine 和 DTR

BaseDet?是基于?MegEngine?的一個(gè)檢測框架，如果要聊 feature，本質(zhì)上也是聊 MegEngine 的 feature，畢竟 BaseDet 只是幫助用戶完成一些基本的訓(xùn)練任務(wù)，有趣的 feature 還是由底層框架支持的，所以這個(gè)部分我們來聊一聊 MegEngine。

為了用戶的遷移性，MegEngine 在一些 API 上和 numpy 做了對齊，這點(diǎn)上和 google 的 jax 是比較類似的，好處是因?yàn)?numpy 的 api 比較穩(wěn)定且 well-known；而 MegEngine 在 module 的上的設(shè)計(jì)比較接近 torch，因?yàn)橛脩魧τ?torch 的 module 的用法是相對熟悉的。對于大部分 torch 用戶，要轉(zhuǎn) MegEngine 還是相對比較絲滑的，最需要注意的點(diǎn)就是：在 MegEngine 里面，autograd 是由一個(gè)叫做 GradManager 的 class 控制的，有點(diǎn)類似 tensorflow 的 GradientTape，這樣做的好處在于方便控制資源的管理，不容易像 torch 一樣出現(xiàn)奇怪的內(nèi)存泄漏現(xiàn)象（對于這個(gè)現(xiàn)象感興趣的同學(xué)，可以參考之前我寫的另一個(gè)?blog：https://zhuanlan.zhihu.com/p/450779978）。

我個(gè)人最喜歡的 MegEngine 的 feature 是由?@圓角騎士魔理沙提（https://www.zhihu.com/people/marisa.moe）出來的 DTR（Dynamic Tensor Rematerialization，推薦去看原文：https://arxiv.org/abs/2006.09616），以 FCOS 的 baseline 為例，在 2080Ti 上單卡訓(xùn)練，不開 DTR batchsize 只能開到 8，打開 DTR 的情況下，batchsize 能翻一倍開到 16（當(dāng)然訓(xùn)練速度也會變慢）。

當(dāng)然，有很多實(shí)現(xiàn)細(xì)節(jié)是原文沒有考慮的，根據(jù) engine 團(tuán)隊(duì)的整理，也在這里分享一些坑點(diǎn)（建議看完論文再來看這里的坑點(diǎn)，理解更深刻一些）：

1. 多卡支持。原始論文沒考慮這個(gè)問題，其實(shí)說起來解決方法很簡單，就是無腦把需要做 send/recv 通訊的 tensor 當(dāng)成 immutable 的，不要 drop 就好了。

2. 顯存碎片經(jīng)常會導(dǎo)致算法不實(shí)用，實(shí)際上估值函數(shù)需要與內(nèi)存分配器聯(lián)動。這里我們?yōu)榱朔奖憷斫馀e個(gè)例子。假設(shè)顯存的狀態(tài)是有 200Mb 可以自由使用，其排布方式是[A(90M) B(10M) C(90M) D(10M)]，其中 A、B、C、D 都是 tensor，括號里面是 tensor 需要的顯存大小。假設(shè)有新的 tensor E 需要 15M 的空間，假設(shè) DTR 默認(rèn)算出來是 drop 掉 B 和 D，但是因?yàn)轱@存不連續(xù)，此時(shí)還需要 drop 掉 A 或者 C，那么一開始 drop 掉 B 和 D的行為就很不劃算，不如一開始就 drop 掉 A 或者 C，所以說估值函數(shù)實(shí)際上是需要和內(nèi)存分配器做聯(lián)動。在 pytorch 里面很難獲取到現(xiàn)在各個(gè) blob 的申請情況，而 mge 里的顯存分配器設(shè)計(jì)的比較干凈，申請釋放也都有統(tǒng)一的地方，所以 DTR 這個(gè)機(jī)制實(shí)現(xiàn)的也相對干凈一些。

3. 涉及跨 iter 操作的時(shí)候會有一些麻煩，比如 ema 中需要進(jìn)行特殊處理，可以參考 BaseDet 里面的示例?code（https://github.com/megvii-research/basedet/blob/0.3.0rc/basedet/layers/common/ema.py#L78-L81）。出現(xiàn)問題的原因在于：諸如 ema 這樣的操作，通常會使得 tensor 的計(jì)算歷史成為一個(gè)無限長（和訓(xùn)練長度一樣）的東西，而 DTR 就會把歷史上用到的 tensor 都記下來（重算過程需要使用），這就會導(dǎo)致出現(xiàn)泄漏現(xiàn)象。

4. 原始論文里收到 pytorch 限需要手動填閾值，大部分用戶并不是很喜歡這種調(diào)用方式，最后在 MegEngine 里面使用的是一個(gè)自適應(yīng)的閾值。對于用戶來說，只需要在 code 里面加上mge.dtr()就能簡單開啟功能了。

不同用戶的不同需求

在曠視內(nèi)部有一個(gè)很棒的帖子，講的是用戶通常只會用到軟件中 15% 的功能，而不同類型的用戶使用的往往是同一個(gè)軟件中那不同的 15% 部分。在完成 BaseDet 的過程中，我接觸到了不同的用戶人群，了解到這些人群對于框架的不同需求。舉個(gè)例子：

• 研究人員：靈活，但同時(shí)有需要的功能的時(shí)候可以簡單打開（比如 ema）。喜歡 pytorch-lightning/timm 這種 lite 的東西，關(guān)心訓(xùn)練/評測邏輯，訓(xùn)練出來的模型點(diǎn)數(shù)越高越好。

• 產(chǎn)品研發(fā)：關(guān)心的重要的參數(shù)能夠簡單配置，方便交付。喜歡 onnx/torchscript 這種中間產(chǎn)物，不關(guān)心訓(xùn)練評測模型的邏輯，像保姆一樣幫他們搞個(gè) demo 走通流程最好。

• 深度學(xué)習(xí)框架研發(fā)：需要簡單就能跑起來的倉庫，方便追溯問題。上層愛咋寫咋寫，愛咋封裝咋封裝，喜歡訓(xùn)練框架提供諸如保存 crashing context、profiler、benchmark 等功能。

所以諸如 mmdet/detectron2 這類框架都是支持簡單的 yaml config 和 lazy eval 的功能的，看起來可能有些矛盾，但是這種做法能夠滿足不同群體的需求。

前面提到過，做檢測應(yīng)用相關(guān)的人則非常傾向于使用 YOLO 系列的各個(gè)框架，一部分原因就是大部分用戶是直接 clone 下來倉庫直接改 code 的，所以在這些框架中，很少提供諸如 registry 和 hook 這類概念，因?yàn)檫@些概念本身并沒有提供靈活性，反而引入了多余的概念。

因?yàn)?BaseDet?本身是為了輔助產(chǎn)品而存在的，所以是基于 product first 的原則而設(shè)計(jì)開發(fā)的，也就不可避免地在使用體驗(yàn)上存在一些 bias，開源出來的目的其實(shí)就是為了糾正這種 bias，還能給 MegEngine 的用戶提供一種code 參考，希望社區(qū)能夠給予一些適當(dāng)?shù)姆答仯@些反饋也是 codebase 前進(jìn)的方向。

BaseDet Github：https://github.com/megvii-research/basedet

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

后記

留下來一段話，送給這世界上愿意花費(fèi)時(shí)間精力去 maintain 項(xiàng)目的開發(fā)人員，也是我這一段時(shí)間來的深刻感悟：任何一段 code 都值得不斷花費(fèi)時(shí)間去打磨，但是打磨之后的 code 并不是真正的產(chǎn)出，關(guān)鍵在于過程中的思考和學(xué)習(xí)。不應(yīng)該和自己維護(hù)的的倉庫過度綁定，總有一些更重要的事情在等著你。

更多 MegEngine 信息獲取，您可以查看：

文檔：https://www.megengine.org.cn/doc/stable/zh/?

GitHub 項(xiàng)目：https://github.com/MegEngine

加入 MegEngine 用戶交流 QQ 群：1029741705

歡迎參與 MegEngine 社區(qū)貢獻(xiàn)，成為?Awesome MegEngineer，榮譽(yù)證書、定制禮品享不停。