mxnet创建新的操作(层)

mxnet创建新的操作(层)

这篇blog将会告诉你如何创建新的MXNet操作(层)。

我们竭尽所能提供最好的操作对于绝大多数的使用场景。然而，如果你发现自己需要自定义层，你有3个选择:

1.使用原生的语言和它的矩阵库(numpy in Python)。这不需要过多的能力和对MXNet的了解。但是他会影响性能。

2.使用原生的语言和mxnet.rtc和mxnet.ndarray。这将给你更好的性能，但是相应的你需要了解更多的MXNet的知识。你可以写CUDA的内核通过Python的方式，并且在运行时进行编译。

3.使用 C++/MShadow(CUDA).这需要你对MXNet，mshadow和Cuda都熟悉。(大神的选择)

Python/Numpy

执行一个python里面的操作和c++里面的操作相似但是更简单了。让我们来创建一个softmax操作，我们通过用mxnet.operator.NumpyOp来开始，然后重写一些方法。

首先，我们调用我们的基础构造器使用need_top_grad=false:

class NumpySoftmax(mx.operator.NumpyOp):
    def __init__(self):
        super(NumpySoftmax, self).__init__(False)

这将告诉引擎我们不需要梯度来自后向传播。

然后我们定义我们的输入和输出

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

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

我们推荐的参数排序方式

['input1', 'input2', ... , 'weight1', 'weight2', ...]

接下去我们需要提供infer_shape来声明我们的output/weight并且检查输入形状的一致性。

    def infer_shape(self, in_shape):
        data_shape = in_shape[0]
        label_shape = (in_shape[0][0],)
        output_shape = in_shape[0]
        return [data_shape, label_shape], [output_shape]

我们的第一个维度总是batch size。我们的标签是一系列的整数，我们的输出和输入有相同的形状大小。Infer_size应该返回两列，即使他们是空的。

 def forward(self, in_data, out_data):
        x = in_data[0]
        y = out_data[0]
        y[:] = np.exp(x - x.max(axis=1).reshape((x.shape[0], 1)))
        y /= y.sum(axis=1).reshape((x.shape[0], 1))

    def backward(self, out_grad, in_data, out_data, in_grad):
        l = in_data[1]
        l = l.reshape((l.size,)).astype(np.int)
        y = out_data[0]
        dx = in_grad[0]
        dx[:] = y
        dx[np.arange(l.shape[0]), l] -= 1.0

记住你需要创建一个新的实例对于每一个symbol。

全部的代码在examples/numpy-ops/numpy_softmax.py

别的两种选择比较麻烦，这里就不讲了。