位址标準化服務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結構	輸入shape	優化前RT	優化後RT	加速比
LSTM - A	7x200	0.199ms	0.066ms	+3.02x
202x200	0.914ms	0.307ms	+2.98x
LSTM - B	70x50	0.266ms	0.098ms	+2.71x
202x50	0.804ms	0.209ms	+3.85x

LSTM單算子優化前後推理性能

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

模型結構	ESIM[6]	ESIM[6]+Blade算子優化	加速比
RT	6.3ms	3.4ms	+1.85x

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倍，意味着可以在提速的同時，讓業務用上更大的模型，獲得更好的業務精度。

模型結構	RT	加速比 （相較 FP32 libtorch）
libtorch	onnxruntime	Blade
位址-2層BERT FP32	17.9ms	17.5ms	14.0ms （編譯優化）	1.28x
位址-4層BERT FP32	37.0ms	34.5ms	28.0ms （編譯優化）	1.32x
位址-4層BERT INT8	21.9ms	26.6ms	9.0ms （量化）	4.11x

位址BERT推理性能展示

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

模型結構	macro F1 （越高越好）
12層BERT-base	77.24
位址-4層BERT	78.72(+1.48)
位址-4層BERT + Blade編譯優化	78.72(+1.48)
位址-4層BERT + Blade量化	78.85(+1.61)

位址BERT相關精度結果

參考文獻

[1] ​​https://help.aliyun.com/document_detail/205129.html​​

[2] ​​https://www.aliyun.com/product/addresspurification/addrp​​

[3] Augmented SBERT: Data Augmentation Method for Improving Bi-Encoders for Pairwise Sentence Scoring Tasks (Thakur et al., NAACL 2021)

[4] Rethink Training of BERT Rerankers in Multi-stage Retrieval Pipeline (Gao et al., ECIR 2021)

[5] ​​https://tianchi.aliyun.com/competition/entrance/531901/introduction​​

[6] Enhanced LSTM for natural language inference[J] (Chen Q, Zhu X, Ling Z, et al., ACL 2017)

[7] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding (Devlin J, Chang M W, Lee K, et al., ACL 2019)

[8] ​​https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html​​