【社區(qū)實(shí)踐】為 TVM 新增 OneFlow 前端
0x0. 背景
去年在Summer Code的時(shí)候我剛好開始入門 TVM（雖然現(xiàn)在仍然也還是入門階段，沒做過什么有意義的工作），并且恰好來到OneFlow 工作就想著給 TVM 添加一個(gè) OneFlow 前端。但可惜在 Summer Code 發(fā)起了這個(gè)項(xiàng)目后因?yàn)橄到y(tǒng)選人的 BUG 導(dǎo)致沒有選到合適的候選人。后來我私下聯(lián)系了申請(qǐng)這個(gè)項(xiàng)目的第二位候選人胡伽魁同學(xué)問他是否愿意來 OneFlow 實(shí)習(xí)并花1-2個(gè)月完成這件事，他同意了并在實(shí)習(xí)期間做了一個(gè)初版出來。感謝胡伽魁同學(xué)的貢獻(xiàn)。
在這個(gè)初版的基礎(chǔ)上，我做了一系列 代碼重構(gòu)，BUG 修復(fù)，文檔編寫，支持更多算子和模型轉(zhuǎn)換之后 使其達(dá)到了一個(gè)相對(duì)穩(wěn)定的狀態(tài)。所以這篇文章來分享一下做這個(gè)小項(xiàng)目的經(jīng)歷和技術(shù)細(xì)節(jié)，希望對(duì)想做開源項(xiàng)目但還沒有做過的讀者提供一個(gè)參考。
0x1. 效果
這里沒有截圖全，可以去官方查看 https://tvm.apache.org/docs/how_to/compile_models/from_oneflow.html 。
Python API預(yù)覽：
現(xiàn)在已經(jīng)成功支持了 ResNet, MobileNet, ShuffleNet，GhostNet，YOLOV3，SRGAN，Vision Transformer在類的多種視覺模型，歡迎大家使用。使用方法見 https://tvm.apache.org/docs/how_to/compile_models/from_oneflow.html 。
0x2. PR歷程
下面的截圖展示了這一工作的 PR 流程，在4月合并了基礎(chǔ)功能的 PR 后基本做的都是 Op 支持和模型支持以及 BUG 修復(fù)。
十分感謝TVM社區(qū)的 「@masahi」 在 PR 過程中的熱心幫助。
0x3. 技術(shù)細(xì)節(jié)
實(shí)際上并沒有什么細(xì)節(jié)可講，基本上就是將OneFlow的IR進(jìn)行逐一遍歷以及逐 Op 轉(zhuǎn)換。我之前已經(jīng)介紹過 TVM 的 ONNX 前端的技術(shù)細(xì)節(jié)了：【從零開始學(xué)TVM】三，基于ONNX模型結(jié)構(gòu)了解TVM的前端 ?，所以這里就不再重復(fù)類似的細(xì)節(jié)了。我這里只列舉一下 OneFlow 前端實(shí)現(xiàn)中的一些特殊一點(diǎn)的細(xì)節(jié)。
形狀和類型推導(dǎo)：對(duì)輸入 Tensor 進(jìn)行形狀和類型推導(dǎo)，功能由 TVM 提供。代碼見：https://github.com/apache/tvm/blob/main/python/tvm/relay/frontend/common.py#L524-L532 形狀和類型提升：對(duì)于如 concat 之類的 Op 來說，如果輸入 Tensor 是不同類型或者不同形狀并且符合提升原則的，那么就可以將其提升到最高類型或者固定的形狀然后再轉(zhuǎn)到 Relay IR。具體實(shí)現(xiàn)見：https://github.com/apache/tvm/blob/main/python/tvm/relay/frontend/oneflow.py#L95-L112 消除獲取 OneFlow Op 的輸入 Tensor 名字的隨機(jī)性：這個(gè)問題是因?yàn)?OneFlow 的 IR 由 Protobuf 做的序列化，所以導(dǎo)致遍歷某個(gè) Node 的時(shí)候拿到輸入的名字是隨機(jī)的，可能會(huì)造成 BUG 。為了解決這一問題，在獲取名字的時(shí)候維護(hù)了一個(gè)有序的列表。具體實(shí)現(xiàn)在：https://github.com/apache/tvm/blob/main/python/tvm/relay/frontend/oneflow.py#L1756-L1765 Relay IR輸入的確定：在 OneFlow IR 中輸入節(jié)點(diǎn)名字都是帶有 _input.這個(gè)特征的，所以根據(jù)這個(gè)特征可以確定 Relay IR 的輸入節(jié)點(diǎn)。具體實(shí)現(xiàn)：https://github.com/apache/tvm/blob/main/python/tvm/relay/frontend/oneflow.py#L1816-L1840
確定了輸入 以及 Op 轉(zhuǎn)換的規(guī)則后，整個(gè) Relay IR 就可以方便的被構(gòu)造出來了，感覺就沒有什么可說的了。如果還想了解更多 TVM 前端具體細(xì)節(jié)就看上面那個(gè)鏈接吧。
0x4. 總結(jié)
本文簡要介紹了筆者和胡伽魁給 TVM 新增 OneFlow 前端的工作，希望對(duì)想做開源項(xiàng)目但還沒有做過的讀者提供一個(gè)參考。
0x5. 參考鏈接
https://github.com/apache/tvm https://github.com/Oneflow-Inc/oneflow