去年在Summer Code的时候我刚好开始入门 TVM(虽然现在仍然也还是入门阶段,没做过什么有意义的工作),并且恰好来到OneFlow 工作就想着给 TVM 添加一个 OneFlow 前端。但可惜在 Summer Code 发起了这个项目后因为系统选人的 BUG 导致没有选到合适的候选人。后来我私下联系了申请这个项目的第二位候选人胡伽魁同学问他是否愿意来 OneFlow 实习并花1-2个月完成这件事,他同意了并在实习期间做了一个初版出来。感谢胡伽魁同学的贡献。
在这个初版的基础上,我做了一系列 代码重构,BUG 修复,文档编写,支持更多算子和模型转换之后 使其达到了一个相对稳定的状态。所以这篇文章来分享一下做这个小项目的经历和技术细节,希望对想做开源项目但还没有做过的读者提供一个参考。
文档预览
这里没有截图全,可以去官方查看 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 。
下面的截图展示了这一工作的 PR 流程,在4月合并了基础功能的 PR 后基本做的都是 Op 支持和模型支持以及 BUG 修复。
PR历程
十分感谢TVM社区的 「@masahi」 在 PR 过程中的热心帮助。
实际上并没有什么细节可讲,基本上就是将OneFlow的IR进行逐一遍历以及逐 Op 转换。我之前已经介绍过 TVM 的 ONNX 前端的技术细节了:【从零开始学TVM】三,基于ONNX模型结构了解TVM的前端 ,所以这里就不再重复类似的细节了。我这里只列举一下 OneFlow 前端实现中的一些特殊一点的细节。
_input.
这个特征的,所以根据这个特征可以确定 Relay IR 的输入节点。具体实现:https://github.com/apache/tvm/blob/main/python/tvm/relay/frontend/oneflow.py#L1816-L1840确定了输入 以及 Op 转换的规则后,整个 Relay IR 就可以方便的被构造出来了,感觉就没有什么可说的了。如果还想了解更多 TVM 前端具体细节就看上面那个链接吧。
本文简要介绍了笔者和胡伽魁给 TVM 新增 OneFlow 前端的工作,希望对想做开源项目但还没有做过的读者提供一个参考。