How to create own operator with python in mxnet?

繼承CustomOp

• 定義操做符，重寫前向後向方法，此時能夠經過_init__ 方法傳遞須要用到的參數

 1 class LossLayer(mxnet.operator.CustomOp):
 2     def __init__(self, *args, **kwargs):
 3         super(LossLayer, self).__init__()
 4         # recipe some arguments for forward or backward calculation
 5         
 6     def forward(self, is_train, req, in_data, out_data, aux):
 7         """
 8         in_data是一個列表，其中tensor的順序和對應屬性類中定義的list_arguments()參數一一對應
 9         out_data輸出列表
10         is_train 是不是訓練過程
11         req [Null, write or inplace, add]指如何處理對應的複製操做
12         """
13         pass
14         # 函數最後通常調用父類的self.assign(dst, req[0], src)進行賦值操做
15         # 但對於dst或者src是list類型的時候要調用屢次assign函數處理，此時也能夠直接本身賦值
16         # dst[:]=src
17         
18     def backward(self, req, out_grad, in_data, out_data, in_grad, aux):
19         """
20         out_grad 上一層反傳的偏差
21         in_data 輸入數據，list
22         out_data 輸出的數據，由forward方法肯定， 其類型大小和out_grad一致
23         in_grad 須要計算的回傳偏差
24         """
25         pass
26         # 其操做值得複製操做相似於forward方法        

• 定義好操做符以後還須要定義其對應的屬性類，並將其註冊到operator中

1 @mx.operator.register('losslayer')  # 注意這裏註冊的名字將是後面調用該操做符使用的類型名

• 重寫對應的屬性類

 1 class LossLayerProp(mx.operator.CustomOpProp): # 這裏的名字並不是必須對應操做類名稱，被@修飾符修飾
 2   def __init__(self, params):
 3     super(LossLayerProp,self).__init__(need_top_grad=False)
 4     # 最後的損失層不須要接收上層的偏差，則將need_top_grad設置爲False
 5     # 能夠傳遞一些參數用以傳遞給操做類
 6    
 7   def list_arguments(self):  
 8     # 這個方法很是重要，定義了該操做符的輸入參數，當綁定對應操做符時，輸入量由該方法指定
 9     return ['data1','data2','data3','label']
10   
11   def list_outputs(self):
12     # 一樣返回的是列表，表示輸出的量,這個實際上是輸出變量的後綴suffix
13     # 若返回的是['output1','output2']則輸出爲 操做類的名稱name加上對應後綴的量[name_output1, name_output2]
14     return ['output']
15   
16   def infer_shape(self, in_shape):
17     # 給定in_shape,顯示每個變量的對應大小，以判斷大小是否一致
18     return [],[],[]
19       # 返回的必須是3個列表，即便列表爲空，分別對應着輸入參數的大小、輸出數據的大小、aux參數的大小，通常最後一個爲空
20     
21     def infer_type(self, in_type):
22       # 該方法相似於infer_shape，推斷數據類型
23 
24     def create_operator(self, ctx, shapes, dtypes):
25       # 該方法真正的建立操做類對象，默認調用
26       return LossLayer()

• 自定義操做符的使用

 1 data1=mx.sym.Variable('data1')
 2 data2=mx.sym.Variable('data2')
 3 data3=mx.sym.Variable('data3')
 4 label = mx.sym.Variable('label')
 5 # 下面這句調用很重要，顯示指定輸入的symbol，而後指定自定義操做符類型
 6 net = mx.sym.Custom(data1=data1, data2=data2, data3=data3, label=label, name='net', op_type='losslayer')  
 7 # 輸出操做符的相關屬性
 8 print(net.infer_shape(data1=(4,1,10,10), data2=(4,1,10,10),data3=(4,1,10,10) label=(4,)))
 9 # data1=(4,1,10,10)表示對應symbol的shape
10 print(net.infer_type(data1=np.int, data2=np.int, data3=np.int, label=np.int))
11 # data1=np.int 標識對應symbol的數據類型
12 print(net.list_arguments()) # 變量參數
13 print(net.list_outputs()) #輸出的變量參數
14 
15 ex = net.simple_bind(ctx=mx.gpu(0), data1=(4,1,10,10), data2=(4,1,10,10),data3=(4,1,10,10) label=(4,)) # simple_bind只須要指定輸入參數的大小
16 ex.forward(data1=data1, data2=data2, label=label))
17 print(ex.outputs[0])

• 上面是沒有參數的層，建立帶有參數的中間層和上面相似， 只是修改下面部分代碼

1 def list_arguments(self):
2     return ['data','weight', 'bias']
3     
4 def infer_shape(self, in_shape):
5     data_shape = in_shape[0]
6     weight_shape = ...
7     bias_shape = ...
8     output_shape = ...
9     return [data_shape, weight_shape, bias_shape], [output_shape], []

調用方式：

net = mx.symbol.Custom(data, name='newLayer', op_type='myLayer')

 包含參數的layer在定義backward方法時要注意梯度的更新方式，即req的選擇

 

NOTE：

有參數的操做符中，通常使用‘weight’和‘bias’做爲參數， 該參數會最爲後綴加到 opname_weight， opname_bias中，由於mxnet默認的參數初始化方法只認‘weight’, 'bias', 'gamma', 'beta'四個量， 對於本身新定義的量，好比weight2， 須要指定初始化方法

Default initialization is now limited to "weight", "bias", "gamma" (1.0), and "beta" (0.0).Please use mx.sym.Variable(init=mx.init.*) to set initialization pattern