Oneflow-Inc / one-fx

A toolkit for developers to simplify the transformation of nn.Module instances. It's now corresponding to Pytorch.fx.
Other
13 stars 1 forks source link

Init the framework (simple). #1

Closed AsakusaRinne closed 1 year ago

AsakusaRinne commented 1 year ago

初始化框架,适配OneFlow0.8.0,目前的计划事项如下:

目前遇到的问题:

  1. OneFlow中缺少ImporterScriptObject,会对某些情况造成影响。
  2. OneFlow中的一些function是直接调到了C API,会导致symbolic trace不能跟踪含有这类function的模块,这里需要想办法处理一下。(已解决)
AsakusaRinne commented 1 year ago

Pytorch中的fx之所以能自动wrap所有的pytorch function是因为pytorch里面有个__pytorch_funtion__的机制,由于one-fx只是个插件,不太能直接照抄,所以对于这个问题目前有两种方案:

  1. 手动列举所有的oneflow function,然后每次都遍历一遍添加wrap(转化成对普通函数的处理)。
  2. 仿照对于nn.module的处理,每次传入函数的时候,判断一下是不是在oneflow的模块里面的(比如'oneflow._oneflow_internal._C'),如果是,那么就走wrap的分支。

目前准备先采用第二种。

AsakusaRinne commented 1 year ago

为了方便后续发包测试和添加文档,先merge掉这个PR,保留commit信息,后续有需要可以回退。