Lightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more
In this tutorial, we will learn how to build custom operators with numpy in python. We will go through two examples:
- Custom operator without any `Parameter`s
- Custom operator with `Parameter`s
Custom operator in python is easy to develop and good for prototyping, but may hurt performance. If you find it to be a bottleneck, please consider moving to a C++ based implementation in the backend.
```python
importnumpyasnp
importmxnetasmx
frommxnetimportgluon,autograd
importos
```
## Parameter-less operators
This operator implements the standard sigmoid activation function. This is only for illustration purposes, in real life you would use the built-in operator `mx.nd.relu`.
### Forward & backward implementation
First we implement the forward and backward computation by sub-classing `mx.operator.CustomOp`:
req : list of {'null', 'write', 'inplace', 'add'}, how to assign to in_grad
out_grad : list of NDArray, gradient w.r.t. output data.
in_grad : list of NDArray, gradient w.r.t. input data. This is the output buffer.
"""
y=out_data[0].asnumpy()
dy=out_grad[0].asnumpy()
dx=dy*(1.0-y)*y
self.assign(in_grad[0],req[0],mx.nd.array(dx))
```
### Register custom operator
Then we need to register the custom op and describe it's properties like input and output shapes so that mxnet can recognize it. This is done by sub-classing `mx.operator.CustomOpProp`:
```python
@mx.operator.register("sigmoid")# register with name "sigmoid"
classSigmoidProp(mx.operator.CustomOpProp):
def__init__(self):
super(SigmoidProp,self).__init__(True)
deflist_arguments(self):
# this can be omitted if you only have 1 input.
return['data']
deflist_outputs(self):
# this can be omitted if you only have 1 output.
return['output']
definfer_shape(self,in_shapes):
"""Calculate output shapes from input shapes. This can be
omited if all your inputs and outputs have the same shape.
in_shapes : list of shape. Shape is described by a tuple of int.
"""
data_shape=in_shapes[0]
output_shape=data_shape
# return 3 lists representing inputs shapes, outputs shapes, and aux data shapes.
return(data_shape,),(output_shape,),()
defcreate_operator(self,ctx,in_shapes,in_dtypes):
# create and return the CustomOp class.
returnSigmoid()
```
### Example Usage
We can now use this operator by calling `mx.nd.Custom`:
```python
x=mx.nd.array([0,1,2,3])
# attach gradient buffer to x for autograd
x.attach_grad()
# forward in a record() section to save computation graph for backward
# see autograd tutorial to learn more.
withautograd.record():
y=mx.nd.Custom(x,op_type='sigmoid')
print(y)
```
```python
# call backward computation
y.backward()
# gradient is now saved to the grad buffer we attached previously
print(x.grad)
```
## Parametrized Operator
In the second use case we implement an operator with learnable weights. We implement the dense (or fully connected) layer that has one input, one output, and two learnable parameters: weight and bias.
The dense operator performs a dot product between data and weight, then add bias to it.
In Linux systems, the default method in multiprocessing to create process is by using fork. If there are unfinished async custom operations when forking, the program will be blocked because of python GIL. Always use sync calls like `wait_to_read` or `waitall` before calling fork.
```python
x=mx.nd.array([0,1,2,3])
y=mx.nd.Custom(x,op_type='sigmoid')
# unfinished async sigmoid operation will cause blocking
os.fork()
```
Correctly handling this will make mxnet depend upon libpython, so the workaround now is to ensure that all custom operations are executed before forking process.
