大家好游昼，我想在這里給大家介紹我的一個項(xiàng)目：MobulaOP.
MobulaOP是一個簡單且靈活的跨框架算子創(chuàng)建工具包。不需要重新編譯深度學(xué)習(xí)框架的源碼场仲，就可以創(chuàng)建自定義的C++算子矛紫。而且只需要一份C++代碼實(shí)現(xiàn)和簡單的定義，自定義算子就可以在CPU和GPU上運(yùn)行揍很。
之所以建立這個項(xiàng)目郎楼，是因?yàn)槲野l(fā)現(xiàn)MXNet創(chuàng)建自定義算子的方法不太方便，其他深度學(xué)習(xí)框架也同樣存在這個問題窒悔。
當(dāng)前呜袁，創(chuàng)建自定義算子的方法主要有：
- 重新編譯深度學(xué)習(xí)框架的源碼
重新編譯源碼耗時(shí)過長。需要了解對應(yīng)框架的算子實(shí)現(xiàn)形式简珠，編寫出的代碼不適用于其他框架阶界。
- 重新編譯深度學(xué)習(xí)框架的源碼
- 使用運(yùn)行時(shí)編譯(Run-Time Compilation)API
需要編寫對應(yīng)的CUDA代碼，編寫過程較復(fù)雜聋庵，無法在CPU環(huán)境下進(jìn)行調(diào)試膘融。
- 使用運(yùn)行時(shí)編譯(Run-Time Compilation)API
- 加載動態(tài)文件
需要了解對應(yīng)框架的動態(tài)加載實(shí)現(xiàn)形式，編寫較復(fù)雜祭玉，一份代碼不適用于多個框架氧映。
- 加載動態(tài)文件
因此，我設(shè)計(jì)了MobulaOP項(xiàng)目攘宙，希望能解決上述問題屯耸。
MobulaOP項(xiàng)目當(dāng)前的特性有：
- 項(xiàng)目實(shí)現(xiàn)精簡，不需要重新編譯深度學(xué)習(xí)框架蹭劈，就可以實(shí)現(xiàn)自定義的C++ operator;
- 只需要編寫一份代碼，就可以讓自定義算子運(yùn)行在不同設(shè)備(CPU/GPU)线召，以及不同的深度學(xué)習(xí)框架(如MXNet, PyTorch)或數(shù)值計(jì)算庫NumPy上铺韧；
- 在編寫自定義層的過程中，用戶有更多的注意力關(guān)注在運(yùn)算的實(shí)現(xiàn)上缓淹；
- 對MXNet有更多的支持哈打，使用MobulaOP可以更方便地創(chuàng)建自定義算子(Custom Operator).
MobulaOP暫時(shí)只支持Linux系統(tǒng)，之后會加入對Windows等系統(tǒng)的支持讯壶。
下面料仗，我想簡單地介紹一下MobulaOP的使用方法。
配置MobulaOP
在終端下輸入以下命令:
# 將MobulaOP項(xiàng)目拷貝下來
git clone https://github.com/wkcn/MobulaOP
# 進(jìn)入項(xiàng)目文件夾
cd MobulaOP
# 安裝依賴庫numpy, pyyaml和easydict
pip install -r requirements.txt
# 進(jìn)行編譯伏蚊，如果需要在GPU下使用立轧，在選項(xiàng)中輸入y
sh build.sh
# 將MobulaOP文件夾加入PYTHONPATH環(huán)境變量中
export PYTHONPATH=$PYTHONPATH:$(pwd)
當(dāng)執(zhí)行完以上命令后，在項(xiàng)目目錄外打開Python交互界面躏吊，輸入import mobula氛改，如果沒有提示，則表示配置成功比伏。
核函數(shù)
配置好MobulaOP后胜卤，就可以使用C++編寫算子(operator)的運(yùn)算函數(shù)了。
這里把并行計(jì)算的運(yùn)算函數(shù)稱為核函數(shù)赁项。
以創(chuàng)建一個逐位乘法算子為例葛躏，它的實(shí)現(xiàn)為：
template <typename T>
MOBULA_KERNEL mul_elemwise_kernel(const int n, const T* a, const T* b, T* out) {
parfor(n, [&](int i) {
out[i] = a[i] * b[i];
});
}
沒錯澈段，定義一個逐位乘法函數(shù)只需要6行代碼，并且它支持在CPU和GPU下運(yùn)行舰攒。
其中败富，MOBULA_KERNEL宏聲明了這個函數(shù)是一個核函數(shù)。核函數(shù)不需要定義返回值芒率，同時(shí)核函數(shù)的函數(shù)名后綴為_kernel.
對于參數(shù)列表囤耳，MobulaOP要求第一個參數(shù)為并行計(jì)算的線程數(shù)。MobulaOP會自動將參數(shù)列表中const T*類型的參數(shù)識別為輸入數(shù)組的指針偶芍，將T*類型的參數(shù)識別為輸出數(shù)組的指針充择。
函數(shù)塊中，調(diào)用了并行執(zhí)行的parfor循環(huán)函數(shù)匪蟀。這個函數(shù)的第一個參數(shù)為循環(huán)體的總迭代數(shù)椎麦，第二個參數(shù)為一個接收迭代下標(biāo)的函數(shù)，這里使用了匿名函數(shù)材彪。下標(biāo)i從0開始計(jì)數(shù)观挎，滿足0 <= i < n。MobulaOP會根據(jù)運(yùn)行設(shè)備對parfor進(jìn)行不同的展開段化。當(dāng)這段代碼在CPU下運(yùn)行時(shí)嘁捷，MobulaOP會將這段函數(shù)展開為：
for (int i = 0; i < n; ++i) {
out[i] = a[i] * b[i];
}
MobulaOP會自動地使用多線程、OpenMP显熏、CUDA等方法并行地執(zhí)行這個循環(huán)雄嚣。
需要注意的是：
-
MOBULA_KERNEL核函數(shù)的第一個參數(shù)為調(diào)用這個函數(shù)進(jìn)行并行計(jì)算的線程數(shù)； - 核函數(shù)內(nèi)部語句均為并行執(zhí)行喘蟆，編寫核函數(shù)時(shí)要注意線程安全問題缓升。當(dāng)前，MobulaOP提供了CPU/GPU下單精度浮點(diǎn)數(shù)(float32)的
atomic_add原子加函數(shù)蕴轨； - 在一個核函數(shù)內(nèi)港谊，允許多次調(diào)用
parfor函數(shù), 這些parfor的總迭代數(shù)可以不同，但實(shí)際使用的線程數(shù)是相同的橙弱； -
parfor函數(shù)只允許在核函數(shù)內(nèi)部進(jìn)行調(diào)用歧寺； - 如果要在核函數(shù)中調(diào)用其他函數(shù)，被調(diào)用的函數(shù)的聲明前需要添加宏
MOBULA_DEVICE, 并聲明返回值類型膘螟。
例子：返回兩個數(shù)中的最大值
template <typename T>
MOBULA_DEVICE T maximum(const T a, const T b) {
return a >= b ? a : b;
}
執(zhí)行核函數(shù)
接下來成福，使用MobulaOP執(zhí)行上述核函數(shù)。
MobulaOP能夠自動分析荆残、生成代碼奴艾，并調(diào)用編譯器將代碼編譯為動態(tài)鏈接庫。
把上述核函數(shù)保存為MulElemWise.cpp文件内斯，放在如下的文件目錄結(jié)構(gòu)中:
tutorial
└── MulElemWise
└─── MulElemWise.cpp
在tutorial文件夾下創(chuàng)建test_mul_func.py文件蕴潦，在這個文件中編寫Python代碼：
import mobula
mobula.op.load('MulElemWise')
import mxnet as mx
a = mx.nd.array([1,2,3])
b = mx.nd.array([4,5,6])
out = mx.nd.empty(a.shape)
mobula.func.mul_elemwise(a.size, a, b, out)
print (out) # [4, 10, 18]
在終端中輸入python test_mul_func.py即可執(zhí)行像啼。
這段代碼中，與MobulaOP相關(guān)的一共有三行(第1潭苞、2忽冻、8行)
第1行代碼導(dǎo)入MobulaOP包。
第2行代碼加載MulElemWise模塊此疹。MobulaOP會搜索MulElemWise文件夾中是否存在同名的.cpp或.py文件僧诚，以及__init__.py文件。若找到這些文件蝗碎，將會對文件進(jìn)行編譯或加載湖笨。mobula.op.load也支持指定搜索目錄，如mobula.op.load('MulElemWise', os.path.dirname(__file__)).
第8行調(diào)用核函數(shù)mul_elemwise蹦骑，與函數(shù)聲明MOBULA_KERNEL mul_elemwise_kernel(const int n, const T* a, const T* b, T* out)相比慈省，在Python中調(diào)用的函數(shù)名比C++中的函數(shù)名少了后綴_kernel. MobulaOP把加載后的核函數(shù)添加到mobula.func中，調(diào)用mobula.func.<核函數(shù)名>即可調(diào)用C++函數(shù)眠菇。MobulaOP能夠自動對參數(shù)進(jìn)行處理, 包括獲取數(shù)據(jù)指針风响、選擇參數(shù)模板沼本、處理內(nèi)存非連續(xù)數(shù)組酣藻、根據(jù)參數(shù)的輸入輸出類型自動調(diào)用wait_to_read藐鹤、wait_to_write函數(shù)等。
創(chuàng)建自定義算子(operator)
如何將核函數(shù)封裝成一個算子(operator)呢登疗，MobulaOP提供了一個簡單的聲明方法怖侦。
在tutorial/MulElemWise文件夾下創(chuàng)建文件MulElemWise.py, 輸入以下代碼：
import mobula
@mobula.op.register
class MulElemWise:
def forward(self, a, b):
mobula.func.mul_elemwise(a.size, a, b, self.y)
def backward(self, dy):
self.dX[0][:] = self.F.multiply(dy, self.X[1])
mobula.func.mul_elemwise(dy.size, dy, self.X[0], self.dX[1])
def infer_shape(self, in_shape):
assert in_shape[0] == in_shape[1]
return in_shape, [in_shape[0]]
第3行的@mobula.op.register為一個Python裝飾器，它將其下面的類注冊為算子谜叹。
一個算子類需要定義forward, backward以及infer_shape函數(shù)。
在forward函數(shù)的參數(shù)列表中搬葬，a和b是算子前向傳播的輸入荷腊；在backward函數(shù)的參數(shù)列表中，dy為算子后向傳播時(shí)輸入的導(dǎo)數(shù)急凰。
MobulaOP會根據(jù)forward函數(shù)得到算子的輸入個數(shù)和名稱女仰，根據(jù)backward得到輸出個數(shù)。
infer_shape函數(shù)傳入的是元組(tuple)的列表抡锈，分別表示各輸入的尺寸(shape). infer_shape的返回值有兩個值疾忍，第一個值是各個輸入的尺寸，第二個值是各個輸出的尺寸床三。infer_shape和MXNet自定義層里的infer_shape是相似的一罩。
在算子的forward和backward函數(shù)中，定義了一些變量：
| 變量名 | 描述 |
|---|---|
| self.F | 當(dāng)前環(huán)境撇簿。假如使用MXNet, self.F = mx.nd |
| self.X[k] | 第k個輸入 |
| self.Y[k] | 第k個輸出 |
| self.dX[k] | 第k個輸入的導(dǎo)數(shù) |
| self.dY[k] | 第k個輸出的導(dǎo)數(shù) |
| self.x | 第1個輸入 |
| self.y | 第1個輸出 |
| self.dx | 第1個輸入的導(dǎo)數(shù) |
| self.dy | 第1個輸出的導(dǎo)數(shù) |
| self.req[k] | 第k個輸入/輸出的處理模式(null/write/add/replace) |
值得注意的是聂渊，當(dāng)使用一個數(shù)組或數(shù)字對另一個數(shù)組賦值時(shí)差购，被賦值的變量后面需要加上[:]，如self.X[0][:] = data
我們也可以使用內(nèi)置的assign函數(shù)進(jìn)行賦值汉嗽，如self.assign(self.X[0], self.req[0], data), 這里的assign函數(shù)和MXNet是一致的欲逃。
測試自定義算子
編寫好MulElemWise算子的定義后，來測試一下吧饼暑。
在tutorial文件夾下創(chuàng)建文件test_mul_op.py, 輸入代碼：
import mobula
mobula.op.load('MulElemWise')
import mxnet as mx
a = mx.nd.array([1,2,3])
b = mx.nd.array([4,5,6])
a.attach_grad()
b.attach_grad()
with mx.autograd.record():
c = mobula.op.MulElemWise(a, b)
c.backward()
print (c) # [4, 10, 18]
print ('a.grad = {}'.format(a.grad.asnumpy())) # [4, 5, 6]
print ('b.grad = {}'.format(b.grad.asnumpy())) # [1, 2, 3]
同樣稳析，在終端輸入python test_mul_op.py指令執(zhí)行。
這里與MobulaOP有關(guān)的新代碼是第11行: c = mobula.op.MulElemWise(a, b)
MobulaOP加載MulElemWise模塊后弓叛，分析了MulElemWise文件夾下的MulElemWise.cpp文件彰居，把核函數(shù)注冊到mobula.func中；同時(shí)加載同一個文件夾下的MulElemWise.py文件邪码，將算子注冊到mobula.op中裕菠。這個過程沒有發(fā)生編譯。
當(dāng)mobula.op.MulElemWise(a, b)執(zhí)行時(shí)闭专，MobulaOP會根據(jù)變量類型奴潘，自動編譯所需要的動態(tài)鏈接庫，并返回結(jié)果影钉。
mobula.op.MulElemWise也可以接受MXNet的符號(Symbol)画髓、Numpy數(shù)組或PyTorch Tensor.
例子：
MXNet的符號(Symbol):
a_sym = mx.sym.Variable('a')
b_sym = mx.sym.Variable('b')
c_sym = mobula.op.MulElemWise(a_sym, b_sym)
Numpy數(shù)組：
a_np = np.array([1,2,3])
b_np = np.array([4,5,6])
# 由于Numpy不支持記錄梯度，因此需要一個實(shí)例記錄梯度
op = mobula.op.MulElemWise[np.ndarray]()
c_np = op(a_np, b_np)
如何在Gluon內(nèi)使用MobulaOP定義的算子呢平委？
我們可以這樣寫：
class MulElemWiseBlock(mx.gluon.nn.HybridBlock):
def hybrid_forward(self, F, a, b):
return mobula.op.MulElemWise(a, b)
這就是MobulaOP的簡單使用介紹奈虾，上述代碼可以在項(xiàng)目的文檔部分(docs)查看。
希望MobulaOP能夠?qū)Υ蠹矣袔椭?/p>
同時(shí)廉赔，歡迎大家對MobulaOP項(xiàng)目提Issue和PR. 謝謝肉微！
