如何創建新的操作符(網絡層)
本節內容描述了創建新的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。