Mxnet学习笔记(3)--自定义Op

https://blog.csdn.net/u011765306/article/details/54562282

前言  

今天因为要用到tile操作(类似np.tile,将数据沿axises进行数据扩充),结果发现mxnet中没有,而且很多操作都没实现,详细完成
度可以参看issue,还在完成中,不过这并不影响我们要用的操作,这里我们
需要实现自己的Op。当然,在官方的example/numpy-ops中已经给出部分例子。这里具体的记录一下。  

自定义Op

自定义op都是去继承operator.py中的类,其中提供如下几类:  

operator.py
CustomOp(object)
CustomOpProp(object)
NDArrayOp(PythonOp)
NumpyOp(PythonOp)
PythonOp(object)
这里很清晰的可以看出,operator分为两条路线,一条路线为CustomOp, 另外一条路线为继承PythonOp,这里我们就分为两部分分别介绍这两条路线。  

CustomOp类  

这条路线是有三步组成,第一步继承CustomOp,重写方法forward()和backward(),然后继承CustomOpProp,重写成员方法,并在方法create_operator中
调用之前写好的Op,第三步调用operator.register()对操作进行注册。具体我们结合官方代码example/numpy-ops/custom_softmax.py来解释,代码如下:  

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))

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))

@mx.operator.register("softmax")
class SoftmaxProp(mx.operator.CustomOpProp):
def __init__(self):
super(SoftmaxProp, self).__init__(need_top_grad=False)

def list_arguments(self):
return ['data', 'label']

def list_outputs(self):
return ['output']

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

def create_operator(self, ctx, shapes, dtypes):
return Softmax()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
上述代码是对softmax的自定义,在类Softmax中重写forward()和backward(),这里与caffe中定义层操作类似,forward()中定义层的前向操作,backward()中
定义反向传播的梯度计算。在完成定义之后,在类SoftmaxProp中create_operator()调用并返回Softmax()实例。那么第三步register如何实现,可以看到,
在SoftmaxProp中带有装饰器mx.operator.register(),等价于操作register("custom_op")(CustomOpProp),这里即在代码运行前即完成了该Op的
实例化,与optimazer的装饰器类似。  

PythonOp类  

这条路线,PythonOp类为基类,而我们大多定义Op时不会去继承它,而是使用他的subclass: NDarrayOp、NumpyOp。这条路线不会像继承CustomOp那样需要三步,这里我们也是只讨论如何继承并定义操作,不去探究
这两个类的实现细节。还是拿官网例子来讲。上代码:  

class NDArraySoftmax(mx.operator.NDArrayOp):
def __init__(self):
super(NDArraySoftmax, self).__init__(False)
self.fwd_kernel = None
self.bwd_kernel = None

def list_arguments(self):
return ['data', 'label']

def list_outputs(self):
return ['output']

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]

def forward(self, in_data, out_data):
x = in_data[0]
y = out_data[0]
if self.fwd_kernel is None:
self.fwd_kernel = mx.rtc('softmax', [('x', x)], [('y', y)])
self.fwd_kernel.push([x], [y], (1, 1, 1), (x.shape[0], 1, 1))

def backward(self, out_grad, in_data, out_data, in_grad):
l = in_data[1]
y = out_data[0]
dx = in_grad[0]
if self.bwd_kernel is None:
self.bwd_kernel = mx.rtc('softmax_grad', [('y', y), ('l', l)], [('dx', dx)])
self.bwd_kernel.push([y,l], [dx], (y.shape[0],1,1), (y.shape[1], 1, 1))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
继承NDArrayOp其实和NumpyOp类似,不同之处在于forward()和backward()重写方式使用函数不同,NDArrayOp中需要使用mx.nd中的操作,而
NumpyOp可以使用numpy中的操作。总之重点在forward()和backward()。当然,如此的自定义方法在使用时需要先定义类对象才可以使用。即与CunstomOp
的定义时间不同。 

成员方法list_arguments,list_outpus,infer_shape

虽然继承方法不同,但是效果是一样的,forward()和backward()是对Op操作的定义,剩余三个成员方法则是对Op接口的描述。  

list_arguments   

该方法主要是对该Op定义时形参的命名,如上述多为['data', 'label'],那么该Op在使用时形参必须为data和label。这里也可以看出mxnet是用过名字
寻找变量的,DataIter,optimazer也是如此。  

list_outputs

同样的,该方法定义了输出变量的名字,一般为opname+’_output’。  

infer_shape

该方法用于在给定输入时,获取该Op的输出shape。当然,在我们自定义时,需要自己设计Op的输入和输出shape。  

以上就是自定义Op时需要做的事情,重点还是forward()和backward(),有时候无头绪的时候可以参考caffe的写法获得灵感。接下来我用例子来说描述一下上述方法。  

import mxnet as mx
import numpy as np
class TileLayer(mx.operator.NumpyOp):
def __init__(self, tiles, axis):
super(TileLayer, self).__init__(False)
# tiles可以为list或者一个数
self.tiles = tiles
self.axis = axis
def list_arguments(self):
return ['input']

def list_outputs(self):
return ['output']

def infer_shape(self, in_shape):
data_shape = in_shape[0]
output_shape = in_shape[0] + [self.tiles]
return [data_shape], [output_shape]

def forward(self, in_data, out_data):
x = in_data[0]
y = out_data[0]
y = np.tile(x, reps=self.tiles)

def backward(self, out_grad, in_data, out_data, in_grad):
bottom_diff = in_grad[0]
top_diff = np.sum(out_grad[0], axis=self.axis)
bottom_diff = top_diff

if __name__ == '__main__':
import logging
from collections import namedtuple
Batch = namedtuple('Batch', ['input'])
logging.basicConfig(level=logging.INFO)
a = mx.sym.Variable('data')
custie = TileLayer(tiles=10, axis=2)
tiles_a = custie(input=a, name='tileop')
arg_shapes, out_shape, aux_shape = tiles_a.infer_shape(data=(2, 3))
logging.info('arg_shape:{} , out_shape:{} , aux_shape:{} , output_blob:{}'.format(arg_shapes, out_shape, aux_shape, tiles_a.list_outputs()))
exe = mx.module.Module(symbol=tiles_a, logger=logging)
exe.bind(data_shapes=[('data', (1, 10, 10))], inputs_need_grad=True)
# exe.init_params()
# exe.init_optimizer()
# data1 = [mx.nd.ones((1, 10, 10))]
# exe.forward(Batch(data1))

# print exe.get_outputs()[0].asnumpy().shape
# top_grads =np.random.random(size=(1, 10, 10, 10))
# exe.backward(out_grads=top_grads)
# print exe.get_input_grads()[0].asnumpy()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
以上为定义的tile操作,这里没有做完全的tile操作,只是可以在最后的axis进行数据的tile操作。forward中用numpy.tile实现,backward中参考caffe
中的TileLayer实现,这里代码运行结果:  

INFO:root:arg_shape:[(2L, 3L)]
out_shape:[(2L, 3L, 10L)]
aux_shape:[]
output_blob:['tileop_output']
1
2
3
4
上述代码因为在list_arguments中定义了形参名字为input,因此在使用是形参必须为input,结果中也可以看到,infer_shape以及list_output的结果,基本细节就是上述。  

在我们定义好Op后,我们需要通过mx.mod.Moudle()将Op进行整合,并通过bind()来申请内存,在此之后,我们可以通过以下两种方法训练它:  

分别调用init_params()初始化参数(当然这里没有参数需要初始化),init_optimazer()初始化optimazer,接下来就可以通过forward()和backward()进行前向反向传播训练模块。  
或者直接调用fit()方法进行训练,因为fit()中包含初始化操作。  
关于Moudle可以参看mx.mod.Module
---------------------
作者:我只是空气
来源:CSDN
原文:https://blog.csdn.net/u011765306/article/details/54562282
版权声明:本文为博主原创文章,转载请附上博文链接!

原文地址:https://www.cnblogs.com/jukan/p/10938380.html