MxNet系列——how_to——new_op

博客新址: http://blog.xuezhisd.top
邮箱：xuezhisd@126.com

---

如何创建新的操作符(网络层)

本节内容描述了创建新的MXNet操作符（或网络）的过程。

我们已经尽了最大努力为最常用的案例提供高性能操作符。然而，如果你需要自定义一个网络层，比如新的损失函数，有两个选择：

• 借助前端语言（比如，Python），使用 CustomOp 来写新的操作符。既可以运行在CPU上，也可以运行在GPU上。根据你的实现情况，性能可能很高，也可能很低。
• 使用 C++/mshadow (CUDA)。如果你不熟悉MXNet，mashadow或Cuda，这是非常困难的，除非你熟悉MXNet，mashadow和Cuda。但是它能获得最佳性能。

CustomOp

在Python中实现一个操作符和在C++中实现方法相似，但更加简单。下面是一个例子，创建了一个softmax操作符。首先构建 mxnet.operator.CustomOp 的子类，然后覆写一些方法：

import os
# 如果自定义的操作符运行在CPU上，MXNET_CPU_WORKER_NTHREADS 必须大于1
os.environ["MXNET_CPU_WORKER_NTHREADS"] = "2"
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))

上面的代码定义了新操作符的前向传播计算函数。前向传播函数的输入包含一个输入（NDArray）的列表和一个输出（NDArray）的列表。为了简便，我们调用输入NDArray的 .asnumpy()，将其转换成基于CPU的NumPy数组。

这样（数据转换）将会非常慢。如果你希望获得最佳性能，那么就保持数据为NDArray格式，并使用mx.nd中的操作来进行计算。

最后，我们使用 CustomOp.assign 将结果数组y赋值给 out_data[0]。它将会根据req的值来决定赋值方式，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 的子类的方式，来定义它的I/O格式。首先，使用名字 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 ，以声明输出和权重的形状，并检查输入形状的一致性：

        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], []

输入/输出张量的第一维是数据批的大小。标签是一组整数，每个整数对应一项数据，并且输出和输入的形状相同。函数 Infer_shape() 应该返回三个列表，即使某一项为空，顺序也必须是：输入，输出和辅助状态（此处没有）。

最后，定义函数 create_operator()。在创建softmax的实例时，后端将会调用该函数：

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

如果想使用自定义操作符，需要创建 mx.sym.Custom 符号。其中，使用新操作符的注册名字来设置参数 op_type：

mlp = mx.symbol.Custom(data=fc3, name='softmax', op_type='softmax')

该例子的完整的代码位于：examples/numpy-ops/custom_softmax.py。

C++/MShadow (CUDA)

更多信息，请查阅 Developer Guide - SimpleOp 和 Developer Guide - Operators。