mxnet創(chuàng)建新的操作(層)

這篇blog將會告訴你如何創(chuàng)建新的MXNet操作(層)。

我們竭盡所能提供最好的操作對于絕大多數的使用場景。然而，如果你發(fā)現自己需要自定義層，你有3個選擇:

1.使用原生的語言和它的矩陣庫(numpy in Python)。這不需要過多的能力和對MXNet的了解。但是他會影響性能。

2.使用原生的語言和mxnet.rtc和mxnet.ndarray。這將給你更好的性能，但是相應的你需要了解更多的MXNet的知識。你可以寫CUDA的內核通過Python的方式，并且在運行時進行編譯。

3.使用 C++/MShadow(CUDA).這需要你對MXNet，mshadow和Cuda都熟悉。(大神的選擇)

Python/Numpy

執(zhí)行一個python里面的操作和c++里面的操作相似但是更簡單了。讓我們來創(chuàng)建一個softmax操作，我們通過用mxnet.operator.NumpyOp來開始，然后重寫一些方法。

首先，我們調用我們的基礎構造器使用need_top_grad=false:

class NumpySoftmax(mx.operator.NumpyOp):def __init__(self):super(NumpySoftmax, self).__init__(False)

這將告訴引擎我們不需要梯度來自后向傳播。

然后我們定義我們的輸入和輸出

def list_arguments(self):return ['data', 'label']def list_outputs(self):return ['output'] 我們推薦的參數排序方式

['input1',?'input2',?...?,?'weight1',?'weight2',?...]

接下去我們需要提供infer_shape來聲明我們的output/weight并且檢查輸入形狀的一致性。

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] 我們的第一個維度總是batch size。我們的標簽是一系列的整數，我們的輸出和輸入有相同的形狀大小。Infer_size應該返回兩列，即使他們是空的。

def forward(self, in_data, out_data):x = in_data[0]y = out_data[0]y[:] = np.exp(x - x.max(axis=1).reshape((x.shape[0], 1)))y /= y.sum(axis=1).reshape((x.shape[0], 1))def backward(self, out_grad, in_data, out_data, in_grad):l = in_data[1]l = l.reshape((l.size,)).astype(np.int)y = out_data[0]dx = in_grad[0]dx[:] = ydx[np.arange(l.shape[0]), l] -= 1.0 記住你需要創(chuàng)建一個新的實例對于每一個symbol。

全部的代碼在examples/numpy-ops/numpy_softmax.py

別的兩種選擇比較麻煩，這里就不講了。

總結

以上是生活随笔為你收集整理的mxnet创建新的操作(层)的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。