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

Posted xiang_freedom

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了MXNet常见问题1怎么创建新运算符(网络层)相关的知识,希望对你有一定的参考价值。

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

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

以上是关于MXNet常见问题1怎么创建新运算符(网络层)的主要内容,如果未能解决你的问题,请参考以下文章

C++ 新运算符。创建新实例

使用标记调度存在自定义 CUDA 分配器时重载类新运算符

奇怪的 C++ 新运算符用法 [重复]

全局关闭 C++ 新运算符异常

用于堆内存分配的新运算符

如何模拟在函数中创建的类(新运算符)?