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

当前的wrap无法做到对间接引用的函数进行封装 #2

Closed AsakusaRinne closed 1 year ago

AsakusaRinne commented 1 year ago

当前的wrap无法做到对间接引用的函数进行封装,比如package a的__init__.py中有一行from b import func,那么当我们在自己的包下执行fx.wrap(a.func)的时候,就会失败。这一点对于oneflow._C这个空间影响比较大,这里面很多常用的函数。

当前的做法是自动wrap所有的oneflow._C空间下的函数,但这样会导致用户也只能用oneflow._C.func的方式去调用,比较不方便。

AsakusaRinne commented 1 year ago

更新:当前已经暂时解决了这个问题,但是可能留有一点小的隐患,那就是按照pytorch中fx的做法,是把类似oneflow.relu这些函数当作call_function来处理的,而现在是当作call_method来处理。这其中原因就在于间接引用的时候frame_dict不全。

这个可能留有一点问题,但当前暂时一切ok。

AsakusaRinne commented 1 year ago

上面这个问题在 https://github.com/Oneflow-Inc/one-fx/pull/1/commits/a7356c48a926c550863059972e3895034cdacd44 这个commit里面做了一下解决,在遇到call_method调用的时候判定一下是不是oneflow中的函数,如果是的话单独处理,避免生成错误代码。

BBuf commented 1 year ago

这样子处理的话,对于真正的Tensor.relu方法,是不是也会被修改成call_function的表示

AsakusaRinne commented 1 year ago

对,确实是也会被修改成call_function的表示,我再试试能不能把这种情况拎出来,之前那样处理主要是有极端情况,比如我们写flow.cat([a, b], axis=1),那么会被生成为[a, b].cat(axis=1),这样就直接跑不通了。但确实有的时候用户可能就是写x.relu这种形式

AsakusaRinne commented 1 year ago

当前主分支已经修改为按照call_function的方式wrap。