地址標準化服務 AI 深度學習模型推理優化實踐

語言: CN / TW / HK

導讀

深度學習已在面向自然語言處理等領域的實際業務場景中廣泛落地,對它的推理效能優化成為了部署環節中重要的一環。推理效能的提升:一方面,可以充分發揮部署硬體的能力,降低使用者響應時間,同時節省成本;另一方面,可以在保持響應時間不變的前提下,使用結構更為複雜的深度學習模型,進而提升業務精度指標。

本文針對地址標準化服務中的深度學習模型開展了推理效能優化工作。通過高效能運算元、量化、編譯優化等優化手段,在精度指標不降低的前提下,AI 模型的模型端到端 推理速度 最高 可獲得了 4.11 倍的提升。

1. 模型推理效能優化方法論

模型推理效能優化是 AI 服務部署時的重要環節之一。一方面,它可以提升模型推理的效率,充分釋放硬體的效能。另一方面,它可以在保持推理延遲不變的前提下,使得業務採用複雜度更高的模型,進而提升精度指標。然而,在實際場景中推理效能優化會遇到一些困難。

1.1 自然語言處理場景優化難點

典型的自然語言處理(Natural Language Processing, NLP)任務中,迴圈神經網路(Recurrent Neural Network, RNN)以及 BERT[7](Bidirectional Encoder Representations from Transformers.)是兩類使用率較高的模型結構。為了便於實現彈性擴縮容機制和線上服務部署的高性價比,自然語言處理任務通常部署於例如 Intel® Xeon® 處理器這樣的 x86 CPU 平臺。然而,隨著業務場景的複雜化,服務的推理計算效能要求越來越高。以上述 RNN 和 BERT 模型為例,其在 CPU 平臺上部署的效能挑戰如下:

  • RNN

迴圈神經網路是一類以序列(sequence)資料為輸入,在序列的演進方向進行遞迴(recursion)且所有節點(迴圈單元)按鏈式連線的遞迴神經網路。實際使用中常見的 RNN 有 LSTM,GRU 以及衍生的一些變種。在計算過程中,如下圖所示,RNN 結構中每一次的後級輸出都依賴於相應的輸入和前級輸出。因此,RNN 可以完成序列型別的任務,近些年在 NLP 甚至是計算機視覺領域被廣泛使用。RNN 相較於與 BERT 而言,計算量更少,模型引數共享,但其計算時序依賴性會導致無法對序列進行平行計算。

RNN 結構示意圖

  • BERT

BERT[7] 證明了能夠以較深的網路結構在大型資料集上完成無監督預訓練(Unsupervised Pre-training),進而供給特定任務進行微調(finetune)的模型。它不僅提升了這些特定任務的精度效能,還簡化了訓練的流程。BERT 的模型結構簡單又易於擴充套件,通過簡單地加深、加寬網路,即可獲得相較於 RNN 結構更好的精度。而另一方面,精度提升是以更大的計算開銷為代價的,BERT 模型中存在著大量的矩陣乘操作,這對於 CPU 而言是一種巨大的挑戰。

BERT 模型結構示意圖

1.2 模型推理優化策略

基於上述推理效能挑戰的分析,我們認為從軟體棧層面進行模型推理優化,主要有如下策略:

  • 模型壓縮:包括量化、稀疏、剪枝等
  • 特定場景的高效能運算元
  • AI 編譯器優化

量化

模型量化是指將浮點啟用值或權重(通常以 32 位元浮點數表示)近似為低位元的整數(16 位元或 8 位元),進而在低位元的表示下完成計算的過程。通常而言,模型量化可以壓縮模型引數,進而降低模型儲存開銷;並且通過降低訪存和有效利用低位元計算指令(如 Intel® Deep Learning Boost Vector Neural Network Instructions,VNNI),取得推理速度的提升。

給定浮點值,我們可以通過如下公式將其對映為低位元值:

其中和是通過量化演算法所得。基於此,以 Gemm 操作為例,假設存在浮點計算流程:

我們可以在低位元域完成相應的計算流程:

高效能運算元

在深度學習框架中,為了保持通用性,同時兼顧各種流程(如訓練),運算元的推理開銷存在著冗餘。而當模型結構確定時,運算元的推理流程僅是原始全量流程個一個子集。因此,當模型結構確定的前提下,我們可以實現高效能推理運算元,對原始模型中的通用運算元進行替換,進而達到提升推理速度的目的。

在 CPU 上實現高效能運算元的關鍵在於減少記憶體訪問和使用更高效的指令集。在原始運算元的計算流程中,一方面存在著大量的中間變數,而這些變數會對記憶體進行大量的讀寫操作,進而拖慢推理的速度。針對這種情況,我們可以修改其計算邏輯,以降低中間變數的開銷;另一方面,運算元內部的一些計算步驟我們可以直接呼叫向量化指令集,對其進行加速,如 Intel® Xeon® 處理器上的高效的 AVX512 指令集。

AI 編譯器優化

隨著深度學習領域的發展,模型的結構、部署的硬體呈現出多樣化演進的趨勢。將模型部署至各硬體平臺時,我們通常會呼叫各硬體廠商推出的 runtime。而在實際業務場景中,這可能會遇到一些挑戰,如:

  • 模型結構、運算元型別的迭代的速度會高於廠家的 runtime,使得一些模型無法快速基於廠商的 runtime 完成部署。此時需要依賴於廠商進行更新,或者利用 plugin 等機制實現缺失的運算元。
  • 業務可能包含多個模型,這些模型可能由多個深度學習框架訓得,此外模型可能需要部署至多個硬體平臺。此時需要將這些格式不同的模型轉化至各個硬體平臺所需的格式,同時要考慮各推理框架實現的不同導致的模型精度效能變化等問題,尤其是像量化這類對於數值差異敏感度較高的方法。

AI 編譯器就是為了解決上述問題而提出的,它抽象出了多個層次來解決上述的一些問題。首先,它接受各個前端框架的模型計算圖作為輸入,並通過各類 Converter 轉化生成統一的中間表示。隨後,諸如運算元融合、迴圈展開等圖優化 pass 會作用至中間表示,以提升推理效能。最後,AI 編譯器會基於優化後的計算圖進行面向特定硬體平臺的 codegen,生成可執行的程式碼,這過程中會引入諸如 stitch、shape constraint 等優化策略。AI 編譯器有很好魯棒性、適應性、易用性,並且能夠收穫顯著優化收益。

本文中,阿里雲機器學習平臺 PAI 團隊聯合英特爾資料中心軟體團隊、英特爾人工智慧和分析團隊、達摩院 NLP 地址標準化團隊,針對地址標準化服務的推理效能挑戰,合作實現了高效能的推理優化方案。

2. 地址標準化介紹

公安政務、電商物流、能源(水電燃)、運營商、新零售、金融、醫療等行業在業務開展的過程中往往涉及大量地址資料,而這些資料往往沒有形成標準結構規範,存在地址缺失、一地多名等問題。隨著數字化的升級,城市地址不標準的問題愈加凸顯。

地址應用現存問題

地址標準化 [2] (Address Purification) 是阿里巴巴達摩院 NLP 團隊依託阿里雲海量的地址語料庫,以及超強的 NLP 演算法實力所沉澱出的高效能及高準確率的標準地址演算法服務。地址標準化產品從規範地址資料、建立統一標準地址庫的角度出發,提供高效能地址演算法。

地址標準化優勢

該地址演算法服務能自動地標準化處理地址資料,可有效地解決一地多名,地址識別,地址真偽辨別等地址資料不規範、人工治理耗時耗力、地址庫重複建設問題,為企業,政府機關以及開發者提供地址資料清洗,地址標準化能力,使地址資料更好的為業務提供支援。地址標準化產品具有如下的幾個特點:

  • 準確率高:擁有海量地址語料庫以及超強的 NLP 演算法技術,並持續優化迭代,地址演算法準確率高
  • 超強效能:積累了豐富的專案建設經驗,能夠穩定承載海量資料
  • 服務全面:提供 20 多種地址服務,滿足不同業務場景需求
  • 部署靈活:支援公共雲、混合雲、私有化部署。

本次優化的模組屬於地址標準化中的搜尋模組。地址搜尋是指使用者輸入地址文字相關資訊,基於地址庫和搜尋引擎,對使用者輸入的地址文字進行搜尋和聯想,並返回相關興趣點(Point of Interest,POI)資訊。地址搜尋功能不僅能夠提升使用者資料處理體驗,同時也是多個地址下游服務的基礎,如經緯度查詢、門址標準化、地址歸一等,因此在整套地址服務體系中起到了關鍵作用。

地址服務搜尋體系示意圖

具體而言,本次優化的模型是基於多工地理預訓練語言模型底座產出的多工向量召回模型精排模型

多工地理預訓練語言模型底座在掩碼語言模型 (Masked Language Model, MLM) 任務的基礎上結合了相關興趣點分類與地址元素識別(省、市、區、POI 等),並通過元學習(Meta Learning)的方式,自適應地調整多個任務的取樣概率,在語言模型中融入通用的地址知識。

多工地址預訓練模型底座示意圖

多工向量召回模型基於上述底座訓練所得,包含雙塔相似度、Geohash (地址編碼) 預測、分詞和 Term Weighting (詞權重) 四個任務。

多工向量召回模型示意圖

作為計算地址相似度匹配的核心模組,精排模型則是在上述底座的基礎上,引入了海量點選資料和標註資料訓練訓練所得 [3],並通過模型蒸餾技術,提升了模型的效率 [4]。最終用應用於召回模型召回的地址庫文件重新排序。基於上述流程訓練得到的 4 層單模型能夠在 CCKS2021 中文 NLP 地址相關性任務 [5] 上獲得較 12 層基線模型更好的效果(詳見效能展示部分)。

精排模型示意圖

3. 模型推理優化解決方案

阿里雲機器學習平臺 PAI 團隊推出的 Blade 產品支援以上提及的所有優化方案,提供了統一的使用者介面,並擁有多個軟體後端,如高效能運算元、Intel Custom Backend、BladeDISC 等等。

Blade 模型推理優化架構圖

3.1 Blade

Blade 是阿里雲機器學習 PAI 團隊(Platform of Artificial Intelligence)推出的通用推理優化工具,可以通過模型系統聯合優化,使模型達到最優推理效能。它有機融合了計算圖優化、Intel® oneDNN 等 vendor 優化庫、BladeDISC 編譯優化、Blade 高效能運算元庫、Costom Backend、Blade 混合精度等多種優化手段。同時,簡潔的使用方式降低了模型優化門檻、提升了使用者體驗和生產效率。

PAI-Blade 支援多種輸入格式,包括 Tensorflow pb、PyTorch torchscript 等。對於待優化的模型,PAI-Blade 會對其進行分析,再應用多種可能的優化手段,並從各種優化結果中選取出加速效果最明顯的為最終的優化結果。

Blade 優化示意圖

為了在保證部署成功率的前提下獲得最大的優化效果,PAI-Blade 採取了 “圈圖” 的方式進行優化,即:

  1. 將待優化子計算圖中,能夠被推理後端 / 高效能運算元支援的部分轉化至相應的優化子圖;
  2. 無法被優化的子圖回退(fallback)至相應的原生框架(TF/Torch)執行。

Blade 圈圖示意圖

Blade Compression 是 Blade 推出的面向模型壓縮的工具包,旨在協助開發人員進行高效的模型壓縮優化工作。它包含了多種模型壓縮功能,包括量化、剪枝、稀疏化等。壓縮後的模型可以便捷地通過 Blade 實現進一步優化,以獲得模型系統聯合的極致優化。

量化方面,Blade Compression:

  • 提供了簡潔的使用介面,通過呼叫幾個簡單 api,即可完成量化改圖、校準(calibration)、量化訓練(Quantization-aware Training,QAT)、匯出量化模型等步驟。
  • 提供了多種後端的支援,通過 config 檔案的配置,即可完成面向不同裝置、不同後端的量化過程。
  • 集成了 PAI-Blade 團隊在實際生產業務中自研的各種演算法,以獲得更高的量化精度

同時,我們提供了豐富的原子能力 api,便於對特定情況進行定製化開發。

Blade Compression 示意圖

BladeDISC 是阿里雲機器學習平臺 PAI 團隊推出的面向機器學習場景的動態 shape 深度學習編譯器,是 Blade 的後端之一。它支援主流的前端框架(TensorFlow、PyTorch)與後端硬體(CPU、GPU),同時也支援推理以及訓練的優化。

BladeDISC 架構圖

3.2 基於 Intel® Xeon® 的 高效能運算元

神經網路模型中的子網路通常具有長期的通用性和普遍性,如 PyTorch 中的 Linear Layer 和 Recurrent Layers 等,是模型建構的基礎模組,負責著特定的功能,通過這些模組的不同組合得到形形色色的模型,並且這些模組也是 AI 編譯器重點優化的目標。據此,為了得到最佳效能的基礎模組,從而實現效能最佳的模型,Intel 針對 X86 架構對這些基礎模組進行了多層次優化,包括使能高效的 AVX512 指令、運算元內部計算排程、運算元融合、快取優化,並行優化等等。

在地址標準化服務中,經常會出現 Recurrent Neural Network (RNN) 模型,並且 RNN 模型中最影響效能的模組是 LSTM 或 GRU 等模組,本章節以 LSTM 為例,呈現在不定長且多 batch 的輸入時,如何實現對 LSTM 的極致效能優化。

通常,為了滿足不同使用者的需求和請求,追求高效能和低成本的雲上服務會將不同的使用者請求進行 Batch,以實現計算資源的最大化利用。如下圖所示,總共有 3 條已經被 embedding 的句子,並且內容和輸入的長度是不相同的。

原始輸入資料

為了使得 LSTM 計算的更高效,需要對 Batched input 採用 PyTorch 的 pack_padded_sequence () 函式進行 padding 和 sort,得到下圖所示,一個 paddding 的資料 tensor,一個描述資料 tensor 的 batch size 的 tensor,一個描述資料 tensor 的原始序號 tensor。

LSTM 輸入資料

到目前為止,已經準備好了 LSTM 的輸入,對於 LSTM 的計算過程如下圖所示,對輸入的 tensor 進行分段批量計算,及跳過零值計算。

LSTM 針對輸入的計算步驟

更深入的 LSTM 的計算優化如下圖 17 所示,公式中的矩陣乘部分進行了公式間計算融合,如下圖所示,原先 4 次矩陣乘轉換成 1 次矩陣乘,並且採用 AVX512 指令進行數值計算,以及多執行緒並行優化,從而實現高效的 LSTM 運算元。其中,數值計算指的是矩陣乘和後序的 elementwise 的元素操作,針對矩陣乘部分,本方案採用的是 oneDNN 庫進行計算,庫中具有高效的 AVX512 GEMM 實現,針對 elementwise 的元素操作,本方案對其採用 AVX512 指令集進行運算元融合,提升了資料在快取中的命中率。

LSTM 計算融合 [8]

3.3 推理後端 Custom Backend

Intel custom backend[9] 作為 Blade 的軟體後端,強有力地加速著模型量化和稀疏的推理效能,主要包含三個層面的優化。首先,採用 Primitive Cache 的策略對記憶體進行優化,其次,進行圖融合優化,最後,在運算元層級,實現了包含稀疏與量化運算元在內的高效運算元庫。

Intel Custom Backend 架構圖

低精度量化

稀疏與量化等高速運算元,得益於 Intel® DL Boost 加速指令集,如 VNNI 指令集。

VNNI 指令介紹

上圖為 VNNI 指令,8bits 可以使用 AVX512 BW 三個指令來加速,VPMADDUBSW 先對 2 對由 8bits 組成的陣列做乘法與加法 , 得到 16bits 資料,VPMADDWD 將相鄰資料加總起來,得到 32bits 資料,最後 VPADDD 加上一個常數,此三函式可組成一個 AVX512_VNNI,此指令可用來加速推理中的矩陣相乘。

圖融合

除此之外,Custom Backend 中也提供了圖融合,例如矩陣相乘後不輸出中間態臨時 Tensor,而是直接執行後面指令,即將後項的 post op 與前級運算元進行融合,如此減少資料搬運以減少執行時間,下圖為一個範例,紅框內的運算元融合後可消除額外的資料搬移,成為一個新的運算元。

圖融合

記憶體優化

記憶體分配與釋放會與作業系統進行通訊,從而導致執行時的延時增加,為了減少這部分的開銷,Custom Backend 中增加了 Primitive Cache 的設計,Primitive Cache 用於快取已經被建立的 Primitive,使得 Primitive 不能被系統回收,減少了下一次呼叫時的建立開銷。

同時為耗時較大的運算元建立了快取機制,以加速運算元執行,如下圖所示:

Primitive Cache

量化功能如之前所說,模型大小減小後,計算與存取的開銷大幅減少,從而效能得到巨大的提升。

4. 整體效能展示

我們選取了地址搜尋服務中典型的兩個模型結構來驗證上述優化方案的效果。測試環境如下所示:

  • 伺服器型號:阿里雲 ecs.g7.large,2 vCPU
  • 測試 CPU 型號:Intel® Xeon® Platinum 8369B CPU @ 2.70GHz
  • 測試 CPU 核數:1 vCPU
  • PyTorch 版本:1.9.0+cpu
  • onnx 版本:1.11.0
  • onnxruntime 版本:1.11.1

4.1 ESIM

ESIM[6] 是一種專為自然語言推斷而生的加強版 LSTM,它的推理開銷主要來自於模型中的 LSTM 結構。Blade 利用 Intel 資料中心軟體團隊開發的高效能通用 LSTM 運算元對其進行加速,替換 PyTorch module 中的預設 LSTM (Baseline)。本次測試的 ESIM 中包含兩種 LSTM 結構,單運算元優化前後的效能如表所示:

LSTM 單運算元優化前後推理效能

優化前後,ESIM 端到端推理速度如表 所示,同時優化前後模型的精度保持不變

ESIM 模型優化前後推理效能

4.2 BERT

BERT[7] 近年來在自然語言處理 (NLP) 、計算機視覺(CV)等領域被廣泛採納。Blade 對該種結構有編譯優化(FP32)、量化(INT8)等多種手段。

速度測試中,測試資料的 shape 固定為 10x53,各種後端及各種優化手段的速度效能如下表所示。可以看到,blade 編譯優化後或 INT8 量化後的模型推理速度均優於 libtorch 與 onnxruntime,其中推理的後端是 Intel Custom Backend & BladeDisc。值得注意的是,經過量化加速後的 4 層 BERT 的速度是 2 層 BERT 的 1.5 倍,意味著可以在提速的同時,讓業務用上更大的模型,獲得更好的業務精度。

地址 BERT 推理效能展示

精度方面,我們基於 CCKS2021 中文 NLP 地址相關性任務 [5] 展示相關模型效能,如下表所示。達摩院地址團隊自研的 4 層 BERT 的 macro F1 精度要高於標準的 12 層 BERT-base。Blade 編譯優化可以做到精度無損,而經過 Blade Compression 量化訓練後的真實量化模型精度要略高於原始浮點模型。

地址 BERT 相關精度結果