MXNet常见问题1怎么创建新运算符(网络层)

Posted 2023-02-15 xiang_freedom

这篇教程教你怎么创建新的运算符或网络层。我们尽可能为大多数情况提供了高效的运算符，然而，如果你是研究人员，你很有可能需要自定义网络层，比如新的损失函数。在这种情况下，你有两个选择：

• 用前端语言（比如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

原文地址：How to Create New Operators (Layers)

推荐：Mxnet学习笔记(3)–自定义Op