Page History

...

The two screenshot are produced with the same code:```

import mxnet as mx
from mxnet import nd
from mxnet import profiler
from mxnet import context
import threading

# Define Custom Operator - Element wise Addition Multiplication
class CustomAddOne(mx.operator.CustomOp):
def forward(self, is_train, req, in_data, out_data, aux):
self.assign(out_data[0], req[0], in_data[0] + 1)

def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
self.assign(in_grad[0], req[0], out_grad[0])

@mx.operator.register("CustomAddOne")
class CustomAddOneProp(mx.operator.CustomOpProp):
def __init__(self):
super(CustomAddOneProp, self).__init__(need_top_grad=True)

def list_arguments(self):
return ['in']

def list_outputs(self):
return ['output']

def infer_shape(self, in_shape):
# inputs, outputs, aux
return [in_shape[0]], [in_shape[0]], []

def create_operator(self, ctx, shapes, dtypes):
return CustomAddOne()

class CustomAddTwo(mx.operator.CustomOp):
def forward(self, is_train, req, in_data, out_data, aux):
y = nd.sqrt(in_data[0])
self.assign(out_data[0], req[0], in_data[0] + 2)

def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
self.assign(in_grad[0], req[0], out_grad[0])

@mx.operator.register("CustomAddTwo")
class CustomAddTwoProp(mx.operator.CustomOpProp):
def __init__(self):
super(CustomAddTwoProp, self).__init__(need_top_grad=True)

def list_arguments(self):
return ['in']

def list_outputs(self):
return ['output']

def infer_shape(self, in_shape):
# inputs, outputs, aux
return [in_shape[0]], [in_shape[0]], []

def create_operator(self, ctx, shapes, dtypes):
return CustomAddTwo()

inp = mx.nd.zeros(shape=(500, 500))
inp2 = mx.nd.zeros(shape=(500, 500))

profiler.set_config(profile_all=True, aggregate_stats=True, continuous_dump = True)
profiler.set_state('run')

w = nd.Custom(inp, op_type="CustomAddOne")
x = nd.sqrt(inp)
y = nd.Custom(inp, op_type="CustomAddOne")
z = nd.Custom(inp, op_type="CustomAddTwo")

mx.nd.waitall()

profiler.set_state('stop')
profiler.dump()
print(profiler.dumps(reset=True))

...

Also, we can try using Gluon:
```

import numpy as np
import mxnet as mx
from mxnet import profiler
from mxnet import gluon, autograd
import os

class Dense(mx.operator.CustomOp):
def __init__(self, bias):
self._bias = bias

def forward(self, is_train, req, in_data, out_data, aux):
x = in_data[0].asnumpy()
weight = in_data[1].asnumpy()
y = x.dot(weight.T) + self._bias
mx.nd.zeros(shape=(500, 500))
self.assign(out_data[0], req[0], mx.nd.array(y))

def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
x = in_data[0].asnumpy()
dy = out_grad[0].asnumpy()
dx = dy.T.dot(x)
self.assign(in_grad[0], req[0], mx.nd.array(dx))

@mx.operator.register("dense") 
class DenseProp(mx.operator.CustomOpProp):
def __init__(self, bias):
super(DenseProp, self).__init__(True)
# we use constant bias here to illustrate how to pass arguments
# to operators. All arguments are in string format so you need
# to convert them back to the type you want.
self._bias = float(bias)

def list_arguments(self):
return ['data', 'weight']

def list_outputs(self):
# this can be omitted if you only have 1 output.
return ['output']

def infer_shape(self, in_shapes):
data_shape = in_shapes[0]
weight_shape = in_shapes[1]
output_shape = (data_shape[0], weight_shape[0])
# return 3 lists representing inputs shapes, outputs shapes, and aux data shapes.
return (data_shape, weight_shape), (output_shape,), ()

def create_operator(self, ctx, in_shapes, in_dtypes):
# create and return the CustomOp class.
return Dense(self._bias)


class DenseBlock(mx.gluon.Block):
def __init__(self, in_channels, channels, bias, **kwargs):
super(DenseBlock, self).__init__(**kwargs)
self._bias = bias
self.weight = self.params.get('weight', shape=(channels, in_channels))

def forward(self, x):
ctx = x.context
return mx.nd.Custom(x, self.weight.data(ctx), bias=self._bias, op_type='dense')




dense = DenseBlock(3, 5, 0.1)
dense.initialize()
x = mx.nd.uniform(shape=(4, 3))


profiler.set_config(profile_all=True, aggregate_stats=True, continuous_dump = True)
profiler.set_state('run')


y = dense(x)


mx.nd.waitall()
profiler.set_state('stop')
profiler.dump()
print(profiler.dumps(reset=True))

...