天天看點

【社群實踐】為 TVM 新增 OneFlow 前端

0x0. 背景

去年在Summer Code的時候我剛好開始入門 TVM(雖然現在仍然也還是入門階段,沒做過什麼有意義的工作),并且恰好來到OneFlow 工作就想着給 TVM 添加一個 OneFlow 前端。但可惜在 Summer Code 發起了這個項目後因為系統選人的 BUG 導緻沒有選到合适的候選人。後來我私下聯系了申請這個項目的第二位候選人胡伽魁同學問他是否願意來 OneFlow 實習并花1-2個月完成這件事,他同意了并在實習期間做了一個初版出來。感謝胡伽魁同學的貢獻。

在這個初版的基礎上,我做了一系列 代碼重構,BUG 修複,文檔編寫,支援更多算子和模型轉換之後 使其達到了一個相對穩定的狀态。是以這篇文章來分享一下做這個小項目的經曆和技術細節,希望對想做開源項目但還沒有做過的讀者提供一個參考。

0x1. 效果

【社群實踐】為 TVM 新增 OneFlow 前端

這裡沒有截圖全,可以去官方檢視 https://tvm.apache.org/docs/how_to/compile_models/from_oneflow.html 。

Python API預覽:

【社群實踐】為 TVM 新增 OneFlow 前端

現在已經成功支援了 ResNet, MobileNet, ShuffleNet,GhostNet,YOLOV3,SRGAN,Vision Transformer在類的多種視覺模型,歡迎大家使用。使用方法見 https://tvm.apache.org/docs/how_to/compile_models/from_oneflow.html 。

0x2. PR曆程

下面的截圖展示了這一工作的 PR 流程,在4月合并了基礎功能的 PR 後基本做的都是 Op 支援和模型支援以及 BUG 修複。

【社群實踐】為 TVM 新增 OneFlow 前端

十分感謝TVM社群的 @masahi 在 PR 過程中的熱心幫助。

0x3. 技術細節

實際上并沒有什麼細節可講,基本上就是将OneFlow的IR進行逐一周遊以及逐 Op 轉換。我之前已經介紹過 TVM 的 ONNX 前端的技術細節了:​​【從零開始學TVM】三,基于ONNX模型結構了解TVM的前端​​ ,是以這裡就不再重複類似的細節了。我這裡隻列舉一下 OneFlow 前端實作中的一些特殊一點的細節。

  • 形狀和類型推導:對輸入 Tensor 進行形狀和類型推導,功能由 TVM 提供。代碼見:https://github.com/apache/tvm/blob/main/python/tvm/relay/frontend/common.py#L524-L532
  • 形狀和類型提升:對于如 concat 之類的 Op 來說,如果輸入 Tensor 是不同類型或者不同形狀并且符合提升原則的,那麼就可以将其提升到最高類型或者固定的形狀然後再轉到 Relay IR。具體實作見:https://github.com/apache/tvm/blob/main/python/tvm/relay/frontend/oneflow.py#L95-L112
  • 消除擷取 OneFlow Op 的輸入 Tensor 名字的随機性:這個問題是因為 OneFlow 的 IR 由 Protobuf 做的序列化,是以導緻周遊某個 Node 的時候拿到輸入的名字是随機的,可能會造成 BUG 。為了解決這一問題,在擷取名字的時候維護了一個有序的清單。具體實作在:https://github.com/apache/tvm/blob/main/python/tvm/relay/frontend/oneflow.py#L1756-L1765
  • Relay IR輸入的确定:在 OneFlow IR 中輸入節點名字都是帶有​

    ​_input.​

    ​這個特征的,是以根據這個特征可以确定 Relay IR 的輸入節點。具體實作:https://github.com/apache/tvm/blob/main/python/tvm/relay/frontend/oneflow.py#L1816-L1840

0x4. 總結

0x5. 參考連結

  • https://github.com/apache/tvm
  • https://github.com/Oneflow-Inc/oneflow