这篇教程教你怎么创建新的运算符或网络层。我们尽可能为大多数情况提供了高效的运算符,然而,如果你是研究人员,你很有可能需要自定义网络层,比如新的损失函数。在这种情况下,你有两个选择:
- 用前端语言(比如python)、使用CustomOp创建新运算符,运行在CPU或者GPU上。不同的实现方式,运行速度可能很快(如果你只用mx.nd下的运算符)或者很慢(如果你使用
.asnumpy()
把数据拷贝出来) - 使用C++/mshadow (CUDA)。这种方式效率最高,但如果你对MXNet,mshadow或者Cuda不熟悉的话,实现起来很难。
CustomOp
用python实现一个运算符很简单。让我们创建一个softmax运算符。继承mxnet.operator.CustomOp
并覆盖几个方法:
import os
import mxnet as mx
import numpy as np
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))
我们定义了运算符的前向计算。forward方法有一系列输入和输出NDArrays。简单起见,我们在输入的第一个NDArray执行.asnumpy()
,把它转化为CPU上执行的Numpy数组。这很慢,如果你想要提高效率,让数据保持NDArray格式并且使用mx.nd下面的运算符。
最后,我们使用 CustomOp.assign把结果数组y赋值给out_data[0],赋值方式取决于req,其取值为‘write’, ‘add’, 或 ‘null’。
对后向传播也是一样:
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))
Softmax定义了我们的自定义运算符的计算,但是你还需要继承mx.operator.CustomOpProp
来定义输入/输出格式。首先,为新运算符注册一个名字‘softmax’:
@mx.operator.register("softmax")
class SoftmaxProp(mx.operator.CustomOpProp):
然后,调用基本构造函数,且设置need_top_grad=False
,因为softmax是损失层,不需要上一层的梯度输入。
def __init__(self):
super(SoftmaxProp, self).__init__(need_top_grad=False)
然后定义输入和输出:
def list_arguments(self):
return ['data', 'label']
def list_outputs(self):
return ['output']
注意list_arguments声明了输入和参数。我们建议以下面的顺序组织:['input1', 'input2', ... , 'weight1', 'weight2', ...]
。
接下来,提供infer_shape
方法声明输出、权重的shape并且检查输入shape的一致性:
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], []
输入、输出张量的第一维代表批数据的不同样本。标签是一列整数,对应每一个数据样本,输出shape和输入shape相同。infer_shape
方法需要返回三个列表:输入,输出和辅助状态(auxiliary states,这里没有),即使有一个是空。
另外,你也可以定义infer_type
来声明输入、输出数据的数据类型。包括np.float32
, np.float64
, np.float16
,np.uint8
和 np.int32
。
def infer_type(self, in_type):
dtype = in_type[0]
return [dtype, dtype], [dtype], []
最后,定义一个create_operator方法,在创建softmax实例的时候会用到:
def create_operator(self, ctx, shapes, dtypes):
return Softmax()
要使用这个自定义运算符,创建一个mx.sym.Custom Symbol,op_type设为注册的名字。
mlp = mx.symbol.Custom(data=fc3, name='softmax', op_type='softmax')
C++
TODO