Mxnet学习笔记(3)--自定义Op

前言　　

今天因为要用到tile操作(类似np.tile，将数据沿axises进行数据扩充)，结果发现mxnet中没有，而且很多操作都没实现，详细完成
度可以参看issue,还在完成中，不过这并不影响我们要用的操作，这里我们
需要实现自己的Op。当然，在官方的example/numpy-ops中已经给出部分例子。这里具体的记录一下。　　

自定义Op

自定义op都是去继承operator.py中的类，其中提供如下几类：　　

• operator.py
  
  • CustomOp(object)
  • CustomOpProp(object)
  • NDArrayOp(PythonOp)
  • NumpyOp(PythonOp)
  • PythonOp(object)

这里很清晰的可以看出，operator分为两条路线，一条路线为CustomOp, 另外一条路线为继承PythonOp,这里我们就分为两部分分别介绍这两条路线。　　

CustomOp类　　

这条路线是有三步组成，第一步继承CustomOp,重写方法forward()和backward(),然后继承CustomOpProp,重写成员方法，并在方法create_operator中
调用之前写好的Op,第三步调用operator.register()对操作进行注册。具体我们结合官方代码example/numpy-ops/custom_softmax.py来解释,代码如下：　　

class Softmax(mx.operator.CustomOp):
    def forward(self, is_train, req, in_data, out_data, aux):
        x = in_data[0].asnumpy()
        y = np.exp(x - x.max(axis=1).reshape((x.shape[0], 1)))
        y /= y.sum(axis=1).reshape((x.shape[0], 1))
        self.assign(out_data[0], req[0], mx.nd.array(y))

    def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
        l = in_data[1].asnumpy().ravel().astype(np.int)
        y = out_data[0].asnumpy()
        y[np.arange(l.shape[0]), l] -= 1.0
        self.assign(in_grad[0], req[0], mx.nd.array(y))

@mx.operator.register("softmax")
class SoftmaxProp(mx.operator.CustomOpProp):
    def __init__(self):
        super(SoftmaxProp, self).__init__(need_top_grad=False)

    def