小米MACE開源項目實踐與思考

何亮亮

人工神經網絡，作為人工智能與機器學習的一個重要分支，最早可以追溯到19世紀40年代。人工神經網絡的歷史，在過去的幾十年中，經歷了起起伏伏。

直到2006年，隨著深度學習這一概念的提出，人工神經網絡的研究再次進入了活躍期，并且引領了這一波人工智能的理論研究和創(chuàng)業(yè)熱潮。這一波深度學習的高速發(fā)展，除了諸如殘差網絡、Drouout等新的理論的提出，更多的是依靠大數(shù)據(jù)和工程上的推動。

首先，互聯(lián)網的發(fā)展、ImageNet等數(shù)據(jù)集以及眾包平臺的建立，為深度學習提供了必需的數(shù)據(jù)。其次，支持通用并行計算的高性能GPU使得深層網絡的訓練成為可能。最后，也是更重要的，是深度學習框架的推出和成熟。

回想一下，在10多年前深度學習復興之前，計算機視覺、自然語言處理等領域的研究人員，要么采用現(xiàn)成的工具，受限于僅能調整少部分元參數(shù)，往往只是通過變換特征提取等方式進行創(chuàng)新，要么為了驗證全新的模型，需要重新造輪子，編寫冗長的C++或者Matlab代碼，研究效率極其低下，這兩者無疑給新理論的創(chuàng)新帶來了巨大的阻礙。

而現(xiàn)在，隨著Theano、Caffe、Torch、TensorFlow等眾多的深度學習框架的推出，算法研究人員，能夠通過編寫簡單的腳本，即可輕松實現(xiàn)復雜的自定義網絡，開發(fā)并驗證自己的模型。

MACE的推出

根據(jù)CBA（Cloud-BigData-AI）戰(zhàn)略，筆者所在的小米人工智能與云平臺部門在2016年開始，逐步加大了對AI領域的投入。AI也逐步賦能公司的多項產品，其中，小米作為一家以手機和IoT智能設備為主的公司，設備端的AI能力是至關重要的。

隨著深度學習領域的發(fā)展，在圖像、NLP和語音領域，傳統(tǒng)方法也逐漸被深度學習方法所取代。

然而，AI的落地過程并非是一帆風順的。成熟的深度學習框架有很多，算法工程師也很容易使用它們開發(fā)出有效的AI模型，但是，如何把訓練好的模型放到設備上高效運行，卻不是一件簡單的事情。

2017年年中，算法團隊在相機的單鏡頭人像模式方面取得了突破性的進展，能夠通過算法，使得單鏡頭能夠達到雙鏡頭的效果，通過后期圖像處理，從而實現(xiàn)媲美單反鏡頭的景深效果。

接下來亟待解決的問題是，如何使該項技術能夠盡快地落地到手機上，這項重任便落在我們團隊身上。

我們首先考察了現(xiàn)有流行的開源框架，諸如TensorFlow、Caffe等，這些框架最初主要面向的是算法研究人員，重點優(yōu)化的方向是服務端的模型訓練，后續(xù)也支持了服務器端的模型推理部署。

所以，最初的設計，以及后續(xù)的優(yōu)化，都是針對服務器環(huán)境，沒有考慮到嵌入式設備內存資源緊缺、計算能力有限，以及功耗敏感等問題。

此外，這些針對服務器端的開源框架，在最重要的算子部分，主要優(yōu)化對象是X86 CPU和英偉達的GPU，對ARM體系結構上的優(yōu)化非常有限。當時，也有一些優(yōu)秀的、專門針對嵌入式ARM優(yōu)化的推理框架，比如國內的ncnn，但也有很大局限，這些框架僅僅對ARM CPU進行了專門優(yōu)化，遠遠不能達到逐幀實時計算的要求。

同時，我們注意到嵌入式設備，尤其是中高端的手機設備，其芯片采用了SoC的形式，整個芯片，除了CPU，還有更強大的GPU和DSP等異構計算單元，而對于圖像視頻類的卷積神經網絡，其模型計算復雜度高，依靠CPU計算還遠遠不能滿足性能指標，如何充分利用這些異構計算單元的算力，是一個成熟的移動端深度學習推理框架必須具備的能力。

開源的框架不能滿足要求，我們把視線轉移到了芯片廠商提供的專有軟件框架上。當時，目標機型是高通的芯片，所以，我們重點考察了SNPE。我們對比了SNPE與TensorFlow和Caffe等開源的框架，發(fā)現(xiàn)SNPE采用GPU或者DSP計算，在速度上大大優(yōu)于這些基于CPU方案的開源框架，并且能夠滿足實時逐幀的計算。

所以，我們決定采用SNPE來部署我們的模型，但隨著項目的推進，我們遇到了一些棘手的問題。

由于當時SNPE發(fā)布不久，還處于早期階段，缺少模型必需的一些算子，同時自定義算子功能不成熟。

另外，SNPE在內存使用上缺少優(yōu)化，沒有考慮不同層之間的緩存復用，導致程序內存占用過大。

同時，由于SNPE采用的是動態(tài)庫的形式，很容易被hook進行反向工程，無法做到安全的模型加密和保護。

為了滿足需求，我們跟高通公司進行了緊密的溝通，但是，由于跨公司溝通環(huán)節(jié)的復雜性，以及看待問題視角和優(yōu)先級的不同，雖然高通給予了我們很多優(yōu)先級的支持，但我們的需求遲遲沒有得到很好的滿足。

隨著時間的推移，項目準時落地的風險越來越大，為了能對底層框架有更好的把控，我們決定在通過變通方式使用SNPE的同時，開發(fā)我們自己的框架，于是MACE（Mobile AI Compute Engine）在2017年8月底正式立項。

最初技術方案有兩個選擇，一個是基于現(xiàn)有開源框架進行修改，另一個是從頭自己實現(xiàn)。

經過分析，我們認為，移動端的推理框架是最核心的異構計算的運行以及底層的算子，而這塊在開源框架中完全缺失。

實際上，框架本身實現(xiàn)難度不高，并且現(xiàn)有的框架非常臃腫，為了能夠適應移動端，需要進行大量的刪減，反而不如重新實現(xiàn)容易。

最終，我們決定采用第二種方式。接下來的一個月，項目進展非常順利，代碼框架迅速搭建起來，并且實現(xiàn)了支持NEON的CPU，能夠成功運行MobileNet。但很快現(xiàn)實給了我們巨大的打擊，評測顯示速度只能達到TensorFlow Mobile速度的1/4，當時正好恰逢國慶節(jié)，但團隊心情卻比較低落。很快，我們完成了MACE的GPU運行，并且實現(xiàn)了1×1的卷積算子，速度相比CPU有所提升，但仍然不理想。

由于在CNN網絡里面，卷積操作是決定性的瓶頸，所以，接下來重點工作便是對卷積操作進行優(yōu)化，于是團隊便分頭行動，對CPU和GPU的卷積算子分別優(yōu)化。在GPU方面，由于現(xiàn)有的開源實現(xiàn)很少且性能表現(xiàn)太差，無參考意義。

我們系統(tǒng)鉆研了并行計算優(yōu)化領域的通用優(yōu)化方法，參考了高通的OpenCL優(yōu)化指南，嘗試了不同的內存布局和計算方式，最終使得1×1卷積有了幾倍的提升，達到了SNPE相近的性能。卷積取得性能突破之后，事情就成功了多半，接下來便是算子的完備，很快自研的MACE框架就能成功運行AI人像模型，相對于最開始的版本，取得了近10倍的性能提升，性能基本接近產品的要求。

接著，隨著產品集成和系統(tǒng)測試工作的進行，除了進一步優(yōu)化算子速度，我們還做了幾項效果顯著的改進，包括Winograd卷積優(yōu)化、算子融合、內存復用、模型加密、啟動優(yōu)化、卡頓優(yōu)化等，使得用戶體驗大大提升。

同時，MACE在CPU的性能優(yōu)化上也取得了突破性進展，在我們的模型上，浮點計算速度達到了TensorFlow的2倍。

2017年年底，離產品代碼凍結時間已經很近，同時進行的SNPE集成仍然不順利，經過對比，在AI人像這個項目上，MACE在各項指標上均優(yōu)于SNPE，并且在快速迭代，所以我們最終決定在產品上上線MACE，至此，MACE項目取得了階段性的成功。

隨著第一款產品的落地，MACE項目重心轉移到了功能的完善和業(yè)務支持上，在后續(xù)的幾個月中，MACE增加了DSP的支持、比特量化支持等重要特性。

MACE的開源

在小米自己的AI落地過程中，我們深深體會到，在手機上落地AI應用絕非易事，而小米之外的應用開發(fā)者面臨同樣的問題。

對于一款手機而言，手機廠商提供的功能畢竟是有限的，一個手機的生態(tài)需要廣大的開發(fā)者共同參與，才能為用戶帶來更優(yōu)異的體驗。

我們從其他渠道了解到，一些實力雄厚的開發(fā)者，例如，算法提供商商湯、曠世等企業(yè)，均有自研的推理框架，并且性能優(yōu)異，但是出于商業(yè)利益的考量，他們也不對外共享。

谷歌作為安卓陣營的引領者，推出了Android NNAPI接口，希望能統(tǒng)一移動端的NN接口，然而，這一努力進展緩慢、無法落地，而谷歌在TensorFlow以及TensorFlow Lite上面，對異構硬件的支持，進展遲緩，遠遠落后于蘋果的生態(tài)。而芯片廠商，諸如高通和聯(lián)發(fā)科，由于各自的商業(yè)考量，技術方案比較封閉，對普通App開發(fā)者不夠友好。

出于這些原因，開發(fā)者面臨的現(xiàn)狀就是只能采用一些基于CPU的方案來進行開發(fā)，這在一些計算密集的CNN模型上非常吃力。

同時，公司CEO雷軍和時任人工智能與云平臺總經理的崔寶秋博士對開源非常支持，尤其是崔博士在開源領域有著多年的經驗，在他們的建議下，我們決定讓MACE項目對外開源，這一決定，在現(xiàn)在看來無疑是非常正確的決定。項目在2018年8月底，正式開源后，得到了社區(qū)的廣泛關注。

開源戰(zhàn)略與意義

對于一個企業(yè)，尤其是新興企業(yè)來說，開源有著非常重要的意義。恰當?shù)牟杉{開源軟件，可以大大降低公司的研發(fā)成本，對外貢獻或者參與開源項目，可以提升公司的技術品牌。

以MACE項目來說，具體可以體現(xiàn)在如下幾個方面。

第一，對外貢獻代碼，尤其是自己主導的項目，可以大大提升研發(fā)人員的成就感，從而提升士氣和開發(fā)效率。

第二，軟件開源可以促使程序員寫代碼時更注重代碼質量，這可大大降低項目的生存周期的管理成本。

第三，隨著開源項目的用戶逐漸積累，會有更多的外部開發(fā)者進行不同形式的貢獻，這對兩家企業(yè)都是互惠互利的事情，從本質上來講，這是更大范圍的社會資源共享。

第四，任何一門基礎技術，都會經歷初期探索、發(fā)展繁榮、成熟標準化的過程，而開源將是大大推進這一進程的重要力量。

對未來的展望

目前，在深度學習推理領域，正經歷著由通用硬件到專有AI加速器的變革過程，無論是IT巨頭，還是AI創(chuàng)業(yè)公司，技術方案百花齊放，各種硬件層出不窮，相應的框架也是眼花繚亂。

這雖然是行業(yè)蓬勃發(fā)展的體現(xiàn)，但由于軟硬件接口千差萬別，給開發(fā)者帶來了很大的障礙。

一方面，開發(fā)者尤其是智能硬件的開發(fā)者，很難對不同的廠家進行客觀的評估，另一方面，工具鏈軟件接口的不同，帶來了額外的開發(fā)成本。

小米作為一家手機和智能硬件廠商，也需要經常評估使用不同的硬件方案，為了簡化評測流程，減少重復勞動，我們開發(fā)了AI評測框架Mobile-AI-Bench項目，并將其開源，我們希望能夠連同社區(qū)的力量，共同將軟硬件評測客觀透明化。

同時，我們也看到，針對硬件碎片化的局面，不同廠商也正在做不同的嘗試。比如，谷歌推出的Android NNAPI試圖在Android平臺上統(tǒng)一硬件接口。ARM也做出了類似的努力，旨在在ARM平臺上統(tǒng)一接口。在MACE項目中，我們也在增加不同加速硬件的支持。

我們相信，在不久的將來，在這一領域，也會逐漸成熟，并標準化，AI計算能力最終會成為一種不可或缺的基礎設施。