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

導讀

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

本文針對地址標準化服務中的深度學習模型開展了推理性能優化工作。通過高性能算子、量化、編譯優化等優化手段，在精度指標不降低的前提下，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 相關精度結果