BaseDet: 走過開發的彎路

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

收到 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 = https://www.cnblogs.com/megengine/p/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，榮譽證書、定制禮品享不停，

標籤：其他

上一篇：微服務組件-----Spring Cloud Alibaba 注冊中心Nacos的CP架構Raft協議分析

下一篇：前端：說說作業中解決過的印象比較深刻的問題