MxNet系列——how_to——new_op

博客新址: http://blog.xuezhisd.top
郵箱：[email protected]

---

如何創建新的操作符(網絡層)

本節內容描述了創建新的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。