作者|MNN團隊
出品|阿里巴巴新零售淘系技術部
MNN 的誕生源於淘系技術部的一群對技術充滿熱情的同學,在充分的行業調研後認為當時的推理引擎如 TFLite 不足以滿足手機淘寶這樣一個億級用戶與日活的超級 App 。
於是我們從零開始自己搭建了屬於阿里巴巴的推理引擎 MNN 。1年前的這個時候,MNN 在 Github 上開源。它比其他的推理引擎更快更輕量,更符合手機淘寶這樣龐大、複雜的生產部署環境。今年3月份,MNN 的引擎設計與優化理念還獲得了學術界的認可,在 MLSys 2020 上發表了論文,並進行了 oral presentation 。
開源1年以來,獲益於公司內外的用戶反饋和業務推動,MNN 在許多方面都取得了長足的進步:
- 在阿里巴巴集團內部得到廣泛推廣,成為了端上推理引擎的事實標準,覆蓋瞭如手機淘寶、手機天貓、優酷、釘釘、閒魚等20多個 App 。
- 新添了模型訓練的支持,從此 MNN 不再是單純的推理引擎,而是具有推理+訓練能力的深度學習引擎。基於 MNN 的訓練能力,我們可以進行 Quantization-Aware Training(QAT)。在 MobileNet 上,MNN 量化訓練之後的模型準確率幾乎不降。
- 持續投資於異構硬件後端的優化,尤其是利用 ARMv8.2 指令集,獲得了兩倍的性能提升。
- 進一步完善 Python 工具鏈,累計新增超過 150 個接口。
- 開源了應用層開箱即用的解決方案 MNNKit ,包含了人臉跟蹤與檢測、人像分割、手勢識別場景的解決方案。
- 開辦了三期《 MNN 學院》直播(1期, 2期, 3期),增加了與用戶們交流,也獲得了忠粉們的大量高質量反饋意見。
截止到今天,MNN 在開源社區獲得了近 4000 的 Github Stars,這是大家對我們的工作所投的 4000 張認可票,也是鞭策我們完善 MNN 的動力。近日,MNN 發佈了 1.0.0 正式版本。自此,MNN 不再被 Github 貼上 “Pre-release” 的標籤了!相較於0.2.2版本,1.0.0 版本的主要升級在於:模型訓練、異構性能和 Python 工具鏈。下面,我們逐項說明。
模型訓練
▐ 模型構建
MNN支持使用 Express (表達式)接口來構建模型,如下例所示,接口還是比較簡潔明瞭的。模型的構建、訓練和保存具體可以參考說明文檔。
VARP x = inputs[0];
x = conv1->forward(x);
x = _MaxPool(x, {2, 2}, {2, 2});
x = conv2->forward(x);
x = _MaxPool(x, {2, 2}, {2, 2});
x = _Convert(x, NCHW);
x = _Reshape(x, {0, -1});
x = ip1->forward(x);
x = _Relu(x);
x = dropout->forward(x);
x = ip2->forward(x);
x = _Softmax(x, 1);
return {x};
以 MNIST 數據集 + Lenet 網絡為例,一個 epoch 60000 張圖片,一般可達到 97-98% 的準確率。性能上,同款 MBP 上,MNN 比 PyTorch 和 Caffe 都有明顯優勢;而手機上,MNN 也達到了完全可用的性能水準。
▐ 量化訓練
模型量化既可以降低模型大小,又可以利用硬件特性提升推理性能,可謂業務應用必備之選。但美中不足之處在於,模型量化會帶來一定的精度損失 —— 對於精度攸關的項目,就難免要做出艱難的選擇了。
為此,MNN 藉助自身模型訓練能力,實現了模型訓練量化,具體實現可以參考說明文檔。精度和壓縮率方面,我們以 MobileNet V2 為例說明:
注1:訓練和驗證均採用 ImageNet 數據集。訓練採用32為 batchsize,執行100個迭代,即,使用了 3200 張圖片進行訓練;精度驗證則使用了 50000 張圖片。
注2:原始模型為 TensorFlow 官方模型,官方準確率為 71.8%,但因預處理代碼上有細微差別,我們測試原始模型的準確率結果稍高於官方;
可以看出,在實現了 73% 模型尺寸壓縮的情況下,量化模型的精度甚至要稍高於原始模型。
▐ 遷移學習示例
這裡節選 MobileNet V2 的 4 分類遷移學習示例,來說明模型的 Finetune,完整示例請參考文檔。
class MobilenetV2TransferModule : public Module {
public:
MobilenetV2TransferModule(const char* fileName) {
// 讀取原始MobilenetV2模型
auto varMap = Variable::loadMap(fileName);
// MobilenetV2的輸入節點
auto input = Variable::getInputAndOutput(varMap).first.begin()->second;
// MobilenetV2分類層之前的節點,AveragePooling的輸出
auto lastVar = varMap["MobilenetV2/Logits/AvgPool"];
// 初始化一個4分類的全連接層,MNN中可以用卷積來表示全連接層
NN::ConvOption option;
option.channel = {1280, 4};
mLastConv = std::shared_ptr<Module>(NN::Conv(option));
// 初始化內部特徵提取器, 內部提取器設成不需要訓練
mFix.reset(PipelineModule::extract({input}, {lastVar}, false));
// 注意這裡只註冊了我們新初始化的4分類全連接層,那麼訓練時將只更新此4分類全連接層
registerModel({mLastConv});
}
virtual std::vector<VARP> onForward(const std::vector<VARP>& inputs) override {
// 輸入一張圖片,獲得MobilenetV2特徵提取器的輸出
auto pool = mFix->forward(inputs[0]);
// 將上面提取的特徵輸入到新初始化的4分類層進行分類
auto result = _Softmax(_Reshape(_Convert(mLastConv->forward(pool), NCHW), {0, -1}));
return {result};
}
// MobilenetV2特徵提取器,從輸入一直到最後一個AveragePooling
std::shared_ptr<Module> mFix;
// 重新初始化的4分類全連接層
std::shared_ptr<Module> mLastConv;
};
int main(int argc, const char* argv[]) {
std::string trainImagesFolder = argv[2];
std::string trainImagesTxt = argv[3];
std::string testImagesFolder = argv[4];
std::string testImagesTxt = argv[5];
// 讀取模型,並替換最後一層分類層
std::shared_ptr<Module> model(new MobilenetV2TransferModule(argv[1])); // arg1: /path/to/mobilenetV2Model
// 進入訓練環節
MobilenetV2Utils::train(model, 4, 0, trainImagesFolder, trainImagesTxt, testImagesFolder, testImagesTxt);
return 0;
}
異構性能
▐ x86
在 x86 上,我們重點優化了矩陣乘法。在分析過 AVX 和 Arm 向量乘指令差異後,我們修改了 AVX 下的權重矩陣佈局,降低了 I/O 佈局,以充分利用 CPU 算力。
此外,我們允許在支持 FMA 擴展的設備上,啟用擴展,將乘法和加法合為一條指令,以進一步降低指令耗時。
當前,FMA 擴展的啟用開關還放置在 CMakeLists.txt 中,後續會在運行時判別。
綜合兩項優化,x86 上有 30% 左右的性能優化。
▐ ARM64
在 ARM64 上,我們面向中低端設備,調整了矩陣乘法的分塊策略,矩陣中每個元素的均攤 I/O 降低了22%;同時,將數據對齊從32字節調整為64字節,與 ARM 架構 CPU 下場景的 L1 cacheline 匹配;最後,優化了緩存預取。優化結果如下:
▐ ARMv8.2
ARM 在「Bringing Armv8.2Instructions to Android Runtime」一文中,列舉了可以在 Android 運行時中應用的 ARMv8.2 新特性。其中,FP16 extensions和Dot Product可以分別應用於浮點計算加速和量化計算加速。
FP16extensions
亦記作asimdhp(Advanced SIMD Half Precision),是 ARMv8.2 架構的可選擴展。asimdhp可用時,可以使用相關 SIMD 指令實現float16的讀寫計算。float16將float32所需的位數降低了一半,因此在 SIMD 下,可以實現兩倍的併發吞吐,從而優化性能。為此,我們在卷積中,採用[N,C/8,H,W,8]的數據佈局,新增了部分卷積實現,效果如下:
精度上幾乎沒有下降,但是性能足足提升了一倍。搭配上 MNN 轉換工具的--fp16輸出選項,模型大小還能減小一半。一箭雙鵰。
Dot Product
亦記作 asimddp(Advanced SIMD Dot Product),是 ARMv8.2 架構的可選擴展。asimddp 可用時,可以使用 SDOT/UDOT 指令實現 int8/uint8 的點積計算。SDOT/UDOT 指令如上圖所示,一次可以處理兩個 4x4 int8/uint8 數據乘,並累加到 4x1 的 int32/uint32 的寄存器上。這樣強大的硬件加速指令,還是雙發射的。
實戰表現效果也非常明顯,在原先 int8 無法發揮效用的設備上,ARMv8.2 也成功實現了耗時減半。
Python 工具鏈
2019年的綠盟開發者大會上,我們發佈了 MNN 的 Python 前端和 Python 版的轉換、量化、可視化工具。而今,Python 又增加了對 MNN Express (表達式)、模型訓練的封裝,累計新增超過 150 個接口。具體可以參考說明文檔。
依然是前文的 Express 構圖,使用 Python 改寫的版本如下:
class Net(nn.Module):
"""construct a lenet 5 model"""
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.conv(1, 20, [5, 5])
self.conv2 = nn.conv(20, 50, [5, 5])
self.fc1 = nn.linear(800, 500)
self.fc2 = nn.linear(500, 10)
def forward(self, x):
x = F.relu(self.conv1(x))
x = F.max_pool(x, [2, 2], [2, 2])
x = F.relu(self.conv2(x))
x = F.max_pool(x, [2, 2], [2, 2])
x = F.convert(x, F.NCHW)
x = F.reshape(x, [0, -1])
x = F.relu(self.fc1(x))
x = self.fc2(x)
x = F.softmax(x, 1)
return x
對熟悉 Python 的開發者來說,是不是要親切上許多呢?
注:目前 Python Express API 處於 BETA 階段。我們會根據社區和內部的反饋持續改進 Python API ,包含進行 backward incompatible 的改動。
後續計劃
2020年,我們計劃每個季度發佈一個穩定版本。
未來的計劃,主要集中在性能、訓練、NPU 支持和模型壓縮。
性能
性能是 MNN 的立身之本,相信很多朋友選擇 MNN,也主要出於它飆車般的性能。有興趣的朋友,可以去看看 MNN 發表在今年 MLSys 的論文解讀。
CPU 上,移動設備方面,ARMv8.2 將是新手機的主流,上文所展示 2 倍加速比非常誘人,我們會進一步挖掘 ARMv8.2 的優化空間;其他平臺方面, x86 的性能在單機訓練、服務端推理的場景中舉足輕重,會是性能優化的另一個目標。
GPU 上,我們會聚焦 Vulkan—— Android 下一代 GPGPU API 的事實標準。
訓練
MNN 最新擁有的訓練能力已經通過 Express (表達式)接口支持常用模型的訓練、量化、蒸餾,我們會進一步完善訓練能力,添加更多算子和求導的支持,以支持更多的模型。
NPU 支持
NPU 具有超高的性能、超低的能耗,將是未來手機的標配。NPU 的支持,也是許多 MNN 用戶經常在釘釘群裡提出的需求。MNN 在未來的1年,會逐步支持更多的 NPU,請大家拭目以待!
模型壓縮
MNN 目前已經擁有 Post-training quantization 和 Quantization-aware training 的能力。我們會持續投入模型壓縮(如蒸餾,稀疏,剪枝,低比特等),給業界提供更多優秀的、即插即用模型壓縮算法。