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

何亮亮

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

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

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

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

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

MACE的推出

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

MACE的開源

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

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

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

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

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

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

開源戰略與意義

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

以MACE項目來說，具體可以體現在如下幾個方面。

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

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

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

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

對未來的展望

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

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

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

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

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

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