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[:] = ydx[np.arange(l.shape[0]), l] -= 1.0记住你需要创建一个新的实例对于每一个symbol。
全部的代码在examples/numpy-ops/numpy_softmax.py
别的两种选择比较麻烦,这里就不讲了。