wkcn / MobulaOP

A Simple & Flexible Cross Framework Operators Toolkit
MIT License
164 stars 21 forks source link

gluon is supported? #6

Closed z01nl1o02 closed 6 years ago

z01nl1o02 commented 6 years ago

firstly, thanks for the work! make it easy to use mxnet.

my question as following: op created with Mobula may be called by gluon?

wkcn commented 6 years ago

Yes. The backend of Gluon is MXNet.

Here is a simple example using MobulaOP in Gluon.

import mxnet as mx
import mobula_op

@mobula_op.operator.register
class MyFirstOP:
    def forward(self, x, y):
        return x + y
    def backward(self, dy): 
        return [dy, dy]
    def infer_shape(self, in_shape):
        assert in_shape[0] == in_shape[1]
        return in_shape, [in_shape[0]]

# Gluon Block
class AddBlock(mx.gluon.nn.HybridBlock):
    def __init__(self):
        super(AddBlock, self).__init__()
    def hybrid_forward(self, F, x, y):
        return MyFirstOP(x, y)

net = AddBlock()
net.initialize()

a = mx.nd.array([1,2,3]) 
b = mx.nd.array([4,5,6])
c = net(a, b)
print (c) # [5, 7, 9]

MobulaOP supports mx.nd.NDArray and mx.sym.Symbol as inputs, so net.hybridize() is available.