xNN：螞蟻端側深度學習框架
?????♀??編者按：本文是支付寶體驗科技沙龍第 3 期-走進螞蟻端智能技術回顧系列文章，作者是螞蟻集團算法工程師世豪，介紹了 xNN 在端側 AI 優(yōu)化問題抽象及解決思路，以及面對“算力碎片化”為特點的 AI 芯片生態(tài)，思考新一代端智能建?？蚣芩璧暮诵囊?，實現(xiàn)對終端設備算力的精細化挖掘。
我今天分享的主題是 《xNN - 螞蟻端側深度學習框架》，接下來首先會簡單介紹螞蟻端智能的發(fā)展歷程，引出我們的第一代 AI 框架建設，我們稱之為叫輕量化建模。在輕量化建模的基礎上，基于我們對端側 AI 的技術發(fā)展趨勢判斷和衍生出的技術挑戰(zhàn)，引申出我們目前正在探索的第二代框架設計-可伸縮建模。
背景介紹
螞蟻端智能訴求可以追溯到 2017 年的支付寶新春掃?；顒?。2017 年是支付寶第一次引入 AR 實景掃福，通過掃描任意“?！弊謳椭蠹胰ゼ占？?。當時的福字識別模型選擇服務端服務部署的技術方案，為了在活動期間識別福字，需要調用大量的服務端部署資源來部署識別模型；另一方面，DL 在云端則意味著數(shù)據(jù)必須上傳。即使不考慮計算壓力，從網(wǎng)絡延時、流量、隱私保護等角度，也給用戶體驗帶來種種限制。因此，對相當多的應用來說，DL 模型前移到移動端部署可以看作是一種剛需。
在這個背景之下，支付寶開始了端側 AI 的探索，并與 2017 年八月在支付寶APP 中首次上線 xNN 端側推理引擎。在同一個時間段前后，我們也看到業(yè)界許多公司都開始布局并落地端上 AI 相關能力，螞蟻也是業(yè)界較早布局端智能技術方向的公司之一。
xNN 引擎上線率先通過圖像識別能力支撐“AR 掃花識花”業(yè)務，并逐步擴展了包括圖像檢測、文字識別、關鍵點回歸等認知智能業(yè)務場景，及以搜索推薦營銷、安全風控為主的數(shù)據(jù)智能業(yè)務場景。同時，端側 AI 能力也在線下 IoT 場景的支付寶刷臉、智能貨柜等中廣泛應用，并通過支付寶小程序、掃一掃等入口將部分成熟的 AI 能力開放到生態(tài)支持。
在廣泛的業(yè)務支持背后，我們需要首先定義端側 AI 的一個核心技術課題?？偨Y來講端側 AI 核心問題是模型優(yōu)化的問題。我們知道在服務端部署 AI 模型時，通常整個部署環(huán)境的資源無論從算力、存儲資源來講，都是相對來講比較豐富的。然而反觀終端移動端設備以及 IOT 設備，無論是算力還是存儲等指標上有若干數(shù)量級下的的嚴苛約束。因此，核心技術挑戰(zhàn)在于如何通過可能的模型優(yōu)化技術手段來降低端上部署模型的性能、尺寸、功耗等等。
作為終端 APP 應用開發(fā)團隊，我們把能夠觸及到的模型優(yōu)化空間抽象為四層：
模型層：通過設計終端友好的模型結構及其相應的參數(shù)來實現(xiàn)在模型任務精度和計算復雜度間的一個權衡。例如目前廣泛應用的 MobileNets、ShuffleNets 系列的模型結構均屬于這個層面上的優(yōu)化。
計算圖層：通過對固化的推理 DAG 計算圖進行算子調度、融合、拆解等手段，保證模型計算依賴關系不變的前提下，提高我們推理效能。社區(qū)和學術界都會有相關的工作。類似 TorchScript、onnx、TVM-Relay 等工作均在這個維度有非常多的能力建設。
算子層：這個維度的優(yōu)化也是端側推理引擎核心關注的部分，即能否通過更為高效的算法、kernel 實現(xiàn)來達到對硬件算力的使用。相關工作開源社區(qū)中各家端側推理引擎，以及服務端 NVIDIA 的 cuDNN 等等加速 Library。
設備層：硬件設備是我們端側推理所依賴的基礎設施，硬件廠商通過提供高效的指令集、AI 專用芯片硬件加速來提高算力峰值。
在此基礎上，螞蟻端側 AI 框架 xNN 定位于覆蓋上述四層優(yōu)化空間。xNN 目標在于構建可以被不同業(yè)務算法在研發(fā)階段復用的端 AI 基礎能力。從模型結構設計、訓練效率優(yōu)化、模型參數(shù)壓縮、端設備的部署，我們希望構建端到端的技術能力來提高整體體驗。在這樣的定位之下，既需要對各個業(yè)務算法領域異同有深入理解、了解其背后的核心訴求，同時也要緊跟整個業(yè)界的 AI 硬件發(fā)展趨勢變化。xNN 團隊也持續(xù)與華為、高通等各廠商保持緊密的合作與溝通。
xNN 框架 1.0：輕量化建模
首先跟大家分享我們在前期所建設的第一代框架的整體設計思路：輕量化建模。
輕量化建模和主流的模型優(yōu)化的思路相對一致。我們基于剛才所定義的四個層級的優(yōu)化主體逐層建設相應的能力，整體上形成一個自上而下的逐層優(yōu)化。在這個過程中，我們沉淀了四個方面的技術能力：結構搜索，模型壓縮，模型轉換，計算引擎。通四個維度逐層向下優(yōu)化，我們希望能夠為我們內部不同的業(yè)務算法需求提供層次化的解決方案。接下來我會對這四個模塊整體做一個簡單的介紹。
首先是模型壓縮和結構搜索，這兩個技術模塊是與端模型算法研發(fā)過程緊密相關，且需要與算法同學深度協(xié)同的一些技術能力。
模型壓縮：
早在 xNN 計算引擎首次上線時，我們便配套布局建設了模型壓縮的核心技術。它包括了模型的剪枝，參數(shù)的量化，以及數(shù)據(jù)編碼。
剪枝（Pruning）：模型剪枝可以簡單的分為三個維度。（1）通道剪枝（Channel Pruning）雖然能夠有效的降低模型的尺寸以及計算的性能，但在實際的使用過程中調整通道間的壓縮比例往往需要大量調參。（2）模版化剪枝（Pattern Pruning）及突觸剪枝（Synapse Pruning）：這種方式通可以在引擎?zhèn)仍O計特殊的稀疏計算 kernel 來實現(xiàn)加速，但往往依賴較高的壓比率才能獲得較好的加速效果。因此想要把這個能力充分的利用起來，往往依賴較大的優(yōu)化精力投入。整體來看，我們認為剪枝相關的技術在逐步的被結構搜索（NAS）的方式所去替代，并逐步將更多的技術資源遷移至 NAS 方向的建設中。
量化（Quantization）：量化是廣泛應用于我們實際的業(yè)務場景中的一項模型壓縮技術。它一方面能夠幫助我們壓縮模型尺寸以減小模型下發(fā)的流量消耗和設備端存儲消耗，比如說我們可以用量化的方式對搜索推薦模型中的 Embedding 的詞表做模型的壓縮；另一方面，我們也可以用量化的方式去做加速。因為很多硬件廠商所提供的指令集都支持了 INT8 量化加速指令。所以量化是一個投入產(chǎn)出性價比較高的模型壓縮手段。
編碼（Coding）：利用數(shù)據(jù)編碼方式來進一步的減小模型的尺寸，是在 xNN 引擎上線第一個模型時便引入的核心能力。我們可以簡單的通過稀疏的表達存儲方式來去做一個信息的存儲；與此同時，xNN 設計了有效的編碼算法，通過對AI的數(shù)據(jù)分布來實現(xiàn)更高的壓縮比率來逼近信息熵的極限。
結構搜索：
網(wǎng)絡結構搜索（Network Architecture Search， NAS）正在 xNN 體系中逐步取代很多單點模型壓縮能力。因為上述的壓縮算法往往會依賴于算法研發(fā)同學優(yōu)化經(jīng)驗和反復嘗試迭代才能獲得較優(yōu)的效果，那么通過自動化的方式來去替代人工的優(yōu)化和搜索，對研發(fā)的效能幫助很大。
在這個背景下，我們也持續(xù)跟進并沉淀了非常多的能力。比如說我們通過 Black-box Optimization 實現(xiàn)自動化的超參搜索，比如通過 Differiential NAS 這種可微的方式，或者是通過 one-shot NAS 的方式來去做模型壓縮。這些方案都廣泛應用在我們的業(yè)務場景中，并取得了很不錯的實踐效果，絕大多數(shù)場景都進一步超過了以往經(jīng)過專家經(jīng)驗深度調優(yōu)的模型指標。
模型轉換：
這個階段主要的工作是把我們主流的深度學習訓練框架所產(chǎn)出的模型轉換為實際部署的 xNN 模型格式。整個的轉換涉及的各種圖優(yōu)化的過程對相關領域同學都比較熟悉，我們不做過多的介紹。
在此基礎之上，圍繞著螞蟻業(yè)務中實際需求，也逐步建設了非常多的增量功能。例如模型串聯(lián)的需求，大家做真正的做端智能的業(yè)務落地，會發(fā)現(xiàn)往往一個業(yè)務所需要的端智能的算法通常并不是由一個單一的模型就能夠解決的，往往是由多個 AI 模型通過一定的算法邏輯串聯(lián)而成。因此有沒有辦法把多個AI模型有效的串聯(lián)到一個模型之中？如果可以，那我們能夠極大的減少算法研發(fā)同學和實際的工程部署同學之間的對接成本，以及整個開發(fā)的生命周期。
因此我們在編譯階段提供了相關的工具：模型串聯(lián)工具幫助把一些包括循環(huán)分支；動態(tài)化工具能夠實現(xiàn)引擎功能、資源調度邏輯的動態(tài)發(fā)布；另外包括模型保護工具、端側訓練工具等等相關的一些能力，都伴隨著業(yè)務的發(fā)展逐步建設和完善起來。
計算引擎：
最后一塊是計算引擎，這個也是我們通常所稱的端側 AI 計算相關的一個核心模塊。計算引擎在我們 APP 的開發(fā)者的視角上來講，我們更多的是聚焦于如何尋找高效的 kernel 實現(xiàn)，在訪存密集型和計算密集型的計算任務上去做合理的優(yōu)化和調整。在整個建設過程中，我們針對不同的部署平臺，會有不同的差異化的一個優(yōu)化策略。比如說在 CPU 和 GPU 上，我們能夠通過自己去實現(xiàn)定制化的高效 kernel 實現(xiàn)，來達到更優(yōu)的效果。因此我們會在這些平臺上去利用我們的專家經(jīng)驗，針對核心算子做深度的優(yōu)化，并會逐步的以更自動化的算子編譯的方式去做演進。
面對現(xiàn)在逐漸發(fā)展起來的 NPU 等 AI 芯片，我們 APP 開發(fā)者往往不太能夠具備非常深入底層的優(yōu)化方式。因此這一塊我們會和各個芯片及硬件廠商會做較多協(xié)同，以期能夠利用最大化的利用到端側 AI 芯片的硬件加速算力。
以上是我們在輕量化建模這個技術建設階段下幾個核心模塊相關的簡單介紹，下面是相關技術在業(yè)務場景中的效果評估，我們主要關注在模型性能、模型尺寸及精度表現(xiàn)。
輕量化建?？偨Y – 優(yōu)化表現(xiàn)
性能：xNN 計算引擎的性能持續(xù)對標業(yè)界多個優(yōu)秀的開源引擎框架，同時我們也圍繞著自己的場景做定制油畫，比如在定點 kernel 配合定點壓縮算法的定制方案能夠在許多定點場景下，能夠達到更具有競爭力的性能表現(xiàn)。
尺寸：模型尺寸極大的影響到了模型下發(fā)的成功率，以及在端側存儲限制。我們在實際的業(yè)務場景中，通過剛剛介紹的各種壓縮優(yōu)化能力，線上模型普遍在百 KB 的量級。
精度：模型壓縮往往會影響原始模型的精度，我們也在持續(xù)關注和并在如量化算法上面的一些技術創(chuàng)新，并持續(xù)保持著在定點優(yōu)化上面的優(yōu)勢。
以上是我們整個對 xNN 輕量化建模階段工作的簡單總結。當然輕量化建模這件事情是需要持續(xù)優(yōu)化投入的課題，我們也會持續(xù)的在不同的核心技術點上去深度打磨。
端智能面臨的問題及挑戰(zhàn)
端 AI 整體發(fā)展趨勢及訴求
在構建 xNN 輕量化建模技術的過程中，我們持續(xù)關注到整個端智能領域的發(fā)展趨勢。我們去看當下的端智能領域，和四五年前去剛開始做端智能的時候相比，已經(jīng)有了一些相當大的變化。我們會把整體端 AI 的一個發(fā)展訴求和趨勢總結為下面三個角度：
從場景角度：我們看到即使是在公司內部已經(jīng)有越來越多樣化的業(yè)務場景希望嘗試端智能來提升整個業(yè)務的一個效果。研發(fā)場景、研發(fā)人數(shù)及模型生產(chǎn)數(shù)量都相比于兩三年前會有數(shù)倍數(shù)的增長。
從 AI 算法角度：這是一個蓬勃發(fā)展和高速演進領域。相比最初的 DNN、CNN 網(wǎng)絡結構，我們看到了如 RNN、LSTM、Transformer 等新的網(wǎng)絡結構設計的演進。這意味著端智能技術側需要持續(xù)的理解以支撐好新的算法架構相關所產(chǎn)生的需求。
從 AI 芯片角度：在四五年前開始做端側 AI 時還是以 CPU 為主的算力支撐。但今天我們再看終端算力環(huán)境，主流的旗艦機型往往都標配了像 NPU 的 AI 硬件加速器。這意味著對于支付寶這樣的國民級 APP，我們需要覆蓋的用戶的手機算力會從 GFlops 到 TFlocps、 3 到 5 個數(shù)量級的變化，算力的碎片化嚴重也意味著我們很難達到一致性的業(yè)務表達。
領域內技術的發(fā)展也帶來短端側 AI 核心訴求的變化。早期我們會投入較多的資源來利用端智能能力實現(xiàn)業(yè)務價值。今天在越來越多的場景和算法研發(fā)已經(jīng)證明了它的價值的基礎之上，我們要考慮的是如何構建我們的技術能力，讓更多的業(yè)務能夠高效的接入，并讓已有的業(yè)務能夠實現(xiàn)更為極致的優(yōu)化。
這其中在關注算法研發(fā)部署效率的基礎之上，一個核心的挑戰(zhàn)是算力的碎片化的問題。如何把我們端上目前已經(jīng)有的更高的算力利用起來。是要去深入去探討的問題。
輕量化建模：算力資源利用率低
我們從算力資源利用率的角度來看一下輕量化建模主要的應用方式。在以往，輕量化建模希望產(chǎn)出極致輕量的一單一的模型。它其實意味著是什么呢？我們要么放棄低端機的用戶覆蓋，要么我們犧牲了高端機用戶的體驗。下面這張圖是一個比較直觀的例子。
我們從 AI benchmark 網(wǎng)站上抽取部分高通平臺隨著年份迭代的芯片算力的增長。我們可以看到如果以注重機型覆蓋為例，使得這個模型在不同的手機上都能夠被它的算力所支撐，達到一個實時的體驗效果。但也就意味著，對于許多新的主流機型，它明明有更高的算力，明明可以用更大模型來達到更好的體驗，但卻不可以做到。
如何最大化 算力資源利用率？
因此我們去考慮如何最大化資源算力的一個利用率呢？那我們考慮如果我們能夠在感知到要部署的具體設備時，能否定制化的夠構建這個設備最優(yōu)的模型。整體來看，我們如果能通過多個模型做部署，就能夠實現(xiàn)對整體算力利用率的提升。這是我們最開始的動機，在這個動機之下，引出了我們目前正在做的第二代 xNN 框架，并稱之為可伸縮建模。
xNN 2.0 可伸縮建模
輕量化建模 -> 可伸縮建模
相比于第一代的框架，我們首先討論下這兩個思路本質的技術區(qū)別是什么。在以往我們追求一個單模型，希望在很多的可能的模型結構中去尋找到一個最優(yōu)解，因此在這個過程中，我們的技術建設思路。始終圍繞著對整個優(yōu)化的路徑上極致提升每一個優(yōu)化工具的功能，來實現(xiàn)一個逐步的逼近最優(yōu)解的方式。
但是如果我們想要把這樣的模式應用在多模型的研發(fā)生產(chǎn)過程中，會存在很多問題。單模型的研發(fā)模式需要大量的人力資源的投入來做端到端的調優(yōu)門檻，如果把這套模式復制到生產(chǎn)多個模型，意味著整個研發(fā)成本會有O（N）的線性復雜度的提升。因此我們希望一種新的研發(fā)部署方式，他能夠有高效的生產(chǎn)多個模型。意味著我們每次的生產(chǎn)過程不是單一的生產(chǎn)最優(yōu)的模型，而是生產(chǎn)一組在不同部署設備上分別最優(yōu)的一個最優(yōu)模型組的集合，即一個在多目標優(yōu)化下的帕雷托前沿模型解幾何。如果我們能夠實現(xiàn)這一步，也就意味著來說我們能夠一次性的產(chǎn)出一批模型，而這些模型能夠分別針對我們所要部署的設備來尋找到它最優(yōu)的模型結構。
基于這樣一個可伸縮建模的思路，我們進一步思考如何生產(chǎn)高質量的、貼近實際生產(chǎn)環(huán)境所需的多模型。我們認為多模型之間它不應該是一個無序的存在，它應該是在某些方面具備了一定的伸縮性的。也就是說我可能在某指標上通過調整某一些核心的因子就能夠實現(xiàn)。比如說我們希望在不同的設備上具備性能上的變化。我調整這些因子就能夠實現(xiàn)不同模型的產(chǎn)出，來達到性能上差異化。只有這些有序的伸縮因子的挖掘，才能有持續(xù)指導我們做更高效的模型研發(fā)部署。所以可伸縮建模其中的核心是可伸縮變換因子的概念。
核心挑戰(zhàn)：算法研發(fā)門檻高
影響多模型指標的核心因子是不是一個相對清晰、確定的答案呢？很明顯它是否定的。學術界的論文往往基于許多理想假設、代理指標來做問題的建模，但是當?shù)剡w移到實際工業(yè)場景中，問題就會變得相當復雜，對于一個模型研發(fā)同學來講的門檻是相當高的。
這里我會舉兩個實際的例子，也是我們在實際的業(yè)務支撐過程中非常有體感的一些例子。比如說剛才我們提到了就是端側 AI 算力的差異化，算力的變化可以通過計算精度的表達來去伸縮。比如說高通的一些新款的設備已經(jīng)能夠支持到 INT4 等混合精度的表達。
如果作為一個算法研發(fā)同學，想要把這個能力應用起來，實現(xiàn)不同性能的表達。那它整個過程中可能會遇到什么阻礙呢？我們會發(fā)現(xiàn)，首先我們需要底層的部署設備具備這樣的一個能力，這本身就是設備差異極大的；同時向上整個的推理引擎?zhèn)纫軌蚍浅：玫募嫒菰O備的能力；同時我們在整個模型編譯的過程中，對設備所依賴的計算圖表達做定制優(yōu)化；我們也依賴定點量化算法上去做一些新的設計；最后我們更需要在整個的模型結構上對設備的功能有相應的理解，才能設計出設備友好的模型結構。也就意味著，如果我想要真正的把推理表達精度的因子廣泛應用起來的話，需要從上到下的一個各個層面上的協(xié)同，只有有充分的協(xié)同，才能夠把這個能力挖掘到最大化。
另一個例子我們從算法的視角來看，上文也提到了整個 AI 算法領域也是在高速的發(fā)展的，很多任務都被 Transformer-based 模型結構刷到了一個更高的一個精度。我們在自己的業(yè)務場景中會發(fā)現(xiàn)如果不針對 Transformer 結構做針對的優(yōu)化，效果是不盡如人意的。那我們再深入的去看這里面為什么呢？例如我們去看 Transformer 里面的 MHSA 核心結構，想要把這個結構給充分的應用起來：在做結構搜索時我們要考慮怎么針對這么多頭注意力機制去搜索可變的參數(shù)；同時思考 MHSA 是否有能夠全定點壓縮算法來達到落地時的性能加速；考慮我們的硬件廠商所提供編譯工具，是不是能夠友好的兼容 Transformer 結構等等；另一方面，如果我們去自己去實現(xiàn)相關 kernel，我們以往的實現(xiàn)有沒有真正關注到這些結構所需要的一些核心算子。只有把這些東西全部的綜合的考慮來看，我們才能得到一個結論，說 Transformer 是不是更適合在某一款具體的 CPU/GPU/NPU 上去執(zhí)行。否則，任何一環(huán)有了缺失，都會影響到最終結論的判斷，進一步影響到算法研發(fā)效果天花板。
因此，如果我們從算法研發(fā)的視角來看，想要讓算法同學去理解這中間的每一環(huán)的領域知識及技術依賴，成本是相當高的。但如果我們能夠以因子的視角把這些影響因子所關聯(lián)的各個層級的技術棧充分的打通，并以一種高效的方式提供給到我們的算法同學，實現(xiàn)面向不同設備的多模型建模，便是我們希望建設的可伸縮建模。
可伸縮建模 – 接入體驗
可伸縮建模，是 xNN 面向下一代建模的建設思路。我們希望做到的是什么呢？就是在以往單模型的算法研發(fā)的代碼基礎之上，僅僅通過少量的代碼修改，就能夠完成多模型建模的遷移。舉例來看，比如說我們想要去構建多個模型以覆蓋從 10ms 到 50ms 的設計目標，我們只需要在原先的模型定義基礎之上，增加一行語法糖去描述 latency 目標的伸縮范圍。接下來，利用我們所沉淀的可伸縮因子、或者是算法同學自研的自己的一些優(yōu)化策略，將原先的靜態(tài)模型結構變換為一個可伸縮的搜索空間描述，最后利用框架的空間-算法解耦設計，使用合適的搜索算法完成最終模型的生產(chǎn)。
以這樣的簡潔的接入體驗，我們就能夠在原先單模型研發(fā)腳本的基礎之上，以一個相對低的接入門檻實現(xiàn)從單模型到多模型的擴展。
總結來看的話，我們探索的可伸縮建模希望圍繞三個目標：
（1）對內我們希望能夠沉淀那些核心的因子來實現(xiàn)跨場景的技術復用；（2）可伸縮建模的框架需要時刻考慮做跨層協(xié)同優(yōu)化的設計，來提升模型端到端優(yōu)化效果的天花板；（3）我們始終要回歸到我們的用戶-算法研發(fā)同學，只有讓算法研發(fā)同學以最低的成本來實現(xiàn)這個目標，我們才能夠把這個技術真正的轉化到我們的業(yè)務之中去。所以就是剛才所提到的整個的研發(fā)界面。我們通過一個更友好的設計來緊緊的透出我們業(yè)務所希望的伸縮的目標和空間，便能夠完成整個建模工作。這是我們希望能夠在未來的一段時間內持續(xù)去建設和優(yōu)化的能力，來提升我們整個業(yè)務的效果。
最后我簡單跟大家做一個總結，我們主要從三個方面分享了 xNN 團隊在技術建設中的思考和沉淀：首先是第一代能力建設輕量化建模中的結構搜索、模型壓縮、模型轉換、推理引擎，它也是我們當下和未來也會持續(xù)去優(yōu)化的；其次，我們基于目前業(yè)務場景、AI 算法高速發(fā)展及端側 AI 算力碎片化問題，定義了新的技術挑戰(zhàn)；最后我們嘗試著去建設我們第二階段的可伸縮建模能力。通過一個低門檻的、設備感知的多模型研發(fā)框架，進一步提升我們對端 AI 算力的挖掘。
