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

語言: CN / TW / HK

0x0. 背景

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

在這個初版的基礎上,我做了一系列 程式碼重構,BUG 修復,文件編寫,支援更多運算元和模型轉換之後 使其達到了一個相對穩定的狀態。所以這篇文章來分享一下做這個小專案的經歷和技術細節,希望對想做開源專案但還沒有做過的讀者提供一個參考。

0x1. 效果

文件預覽

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

Python API預覽:

Python API預覽

現在已經成功支援了 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社群的 「@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

確定了輸入 以及 Op 轉換的規則後,整個 Relay IR 就可以方便的被構造出來了,感覺就沒有什麼可說的了。如果還想了解更多 TVM 前端具體細節就看上面那個連結吧。

0x4. 總結

本文簡要介紹了筆者和胡伽魁給 TVM 新增 OneFlow 前端的工作,希望對想做開源專案但還沒有做過的讀者提供一個參考。

0x5. 參考連結

  • https://github.com/apache/tvm

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