1. 為什么我非要把 HuggingFace 的英譯中模型搬到 ONNX先說說這件事的背景。我手頭有個小項目核心功能是給一批英文技術(shù)文檔做實時翻譯摘要量不大但要求延遲低、部署環(huán)境干凈最好不依賴 GPU 就能跑。最開始我直接用了 HuggingFace 上的Helsinki-NLP/opus-mt-en-zh模型這套 MarianMT 系列的英譯中效果在通用場景下挺穩(wěn)的尤其在處理日常表達和短句時翻譯質(zhì)量比不少在線接口都自然。但問題出在部署環(huán)節(jié)Python 環(huán)境 Transformers 庫 PyTorch 全家桶這套組合在開發(fā)機上跑沒問題一放到生產(chǎn)環(huán)境的容器里就變得非常臃腫啟動時加載模型要等好幾秒單條短句的推理延遲也總是壓不進我想要的閾值內(nèi)。網(wǎng)上不少人提到把模型轉(zhuǎn)成 ONNX說這樣可以擺脫 PyTorch Runtime 的依賴用 ONNX Runtime 直接推理。我當時的第一反應(yīng)是就這么簡單轉(zhuǎn)完真能用實測下來發(fā)現(xiàn)方向是對的但中間坑比想象中多。把這套流程完整走一遍之后我覺得值得把經(jīng)驗整理出來。本文適合兩類人一類是已經(jīng)跑通 HuggingFace 翻譯模型、想優(yōu)化部署體積和延遲的開發(fā)者另一類是剛接觸 ONNX想知道從 PyTorch 到 ONNX這條路上到底有哪些繞不開的細節(jié)的初學(xué)者。先給個結(jié)論Helsinki-NLP/opus-mt-en-zh這類基于 MarianMT 架構(gòu)的模型轉(zhuǎn) ONNX 之后在 CPU 上的推理速度大約能提升 1.5 到 2.5 倍模型體積多多少少會壓縮一些最關(guān)鍵的是運行時可以徹底不裝 PyTorch只保留 onnxruntime 一個推理引擎。但如果你以為 HuggingFace 的from_pretrained加載方式能無縫套用到 ONNX 上那大概率會卡在 tokenizer 和動態(tài)維度這兩個坎上我后面會詳細講。我這次遷移的目標很明確用 ONNX Runtime 替代 Transformers PyTorch 做 CPU 推理保證翻譯質(zhì)量基本不變延遲放到可接受范圍。整體流程圖大概是這樣的選模型 → 驗證原始效果 → 導(dǎo)出 ONNX → 處理動態(tài)軸和 tokenizer 細節(jié) → 量化壓縮 → 用 ONNX Runtime 寫推理腳本 → 對比質(zhì)量和性能?，F(xiàn)在一步步拆開說。2. 首先得搞清楚HuggingFace 里的英譯中模型到底是怎么組織的這一步很重要因為很多人直接跳到torch.onnx.export就動手了結(jié)果被各種張量維度報錯搞得頭大。我建議先花十分鐘把模型的結(jié)構(gòu)和 tokenizer 行為摸清楚后面所有步驟都會順利很多。2.1 選用 opus-mt-en-zh 而不是其他模型的理由HuggingFace 上的英譯中模型不少常見的有Helsinki-NLP/opus-mt-en-zh、facebook/m2m100_418M、t5-small微調(diào)版本等。我最終選了opus-mt-en-zh原因有三第一它是真正的 encoder-decoder 架構(gòu)源語言英文、目標語言中文處理長度適中的句子時效果穩(wěn)定而且不需要像 m2m100 那樣還得額外指定語言代碼 token簡化了預(yù)處理邏輯。第二它的參數(shù)量只有大約 300MB 級別相對于 m2m100 的 418M 甚至更大的模型在 CPU 上推理的壓力小很多。第三這個模型在 HuggingFace 上的下載量和社區(qū)討論度都很高遇到問題容易找到參考這對做遷移的人來說非常重要。當然如果你的場景是長文檔翻譯、或者對特定領(lǐng)域的術(shù)語要求很高可以考慮換用更大的模型。但我這篇博客里的方法和步驟只要架構(gòu)是 encoder-decoder 類的基本都通用。2.2 從 AutoTokenizer 到 tokenizer 的底層行為為什么轉(zhuǎn)換時它最容易出問題跑AutoTokenizer.from_pretrained(Helsinki-NLP/opus-mt-en-zh)之后你拿到的是一個 MarianTokenizer。它和常見的 BERT 類 tokenizer 有個顯著區(qū)別它內(nèi)置了源語言和目標語言的 vocab且附帶一組特殊的控制 token比如/s、pad以及語言標記。在默認調(diào)用tokenizer(text, return_tensorspt)時它會自動完成 padding、truncation并且返回 PyTorch tensor 格式的input_ids和attention_mask。但當你手動導(dǎo)出 ONNX 模型時你通常使用的是 HuggingFace 的transformers.onnx工具包或者直接用torch.onnx.export。問題就出在這里ONNX 模型的輸入要求是確定的張量形狀和數(shù)據(jù)類型而 tokenizer 的__call__方法里有一大堆動態(tài)邏輯padding 到 batch 內(nèi)最大長度、特殊 token 的拼接、mask 的生成等。這些邏輯無法直接映射到 ONNX 的計算圖里。所以正規(guī)做法是把 tokenizer 留在 Python 側(cè)處理ONNX 模型只負責(zé)張量進張量出。也就是說我們用 tokenizer 把文本轉(zhuǎn)成input_ids和attention_mask喂給 ONNX 模型拿到logits之后再回到 Python 側(cè)做 decode。這個分工在后期寫推理腳本時特別重要很多人把 tokenizer 也試圖塞進 ONNX 圖里結(jié)果搞得非常復(fù)雜且毫無必要。2.3 模型的輸入輸出簽名encoder-decoder 的隱藏參數(shù)再細看一下輸入輸出。MarianMT 模型在默認調(diào)用時的輸入是input_ids、attention_mask如果是 decoder 階段還需要decoder_input_ids和encoder_outputs。直接導(dǎo)出整個模型包括 decoder loop非常困難因為內(nèi)部有循環(huán)邏輯ONNX 本身不支持 Python 式的動態(tài)循環(huán)雖然有 Loop 算子但實現(xiàn)復(fù)雜。所以社區(qū)普遍采用的方案是導(dǎo)出兩個 ONNX 模型一個 encoder一個 decoder然后在 Python 側(cè)自己寫解碼循環(huán)說白了就是逐 token 生成。這是絕大多數(shù) HuggingFace ONNX 導(dǎo)出的默認做法transformers.onnx工具也是這樣干的。你會在導(dǎo)出的文件夾里看到encoder_model.onnx和decoder_model.onnx兩個文件原因就在這里。另一個關(guān)鍵參數(shù)是use_cache即 past_key_values。在 PyTorch 推理時model.generate()會自動緩存每一層 attention 的 key/value避免重復(fù)計算。導(dǎo)出 ONNX 時為了做到同樣的效果decoder 模型的輸入里需要顯式聲明past_key_values相關(guān)的輸入。HuggingFace 的 ONNX 導(dǎo)出腳本已經(jīng)處理好了這些細節(jié)所以你在配置OverridableConfig時會看到相關(guān)選項但用現(xiàn)成工具時不需手工干預(yù)。3. 遷移前的準備工作環(huán)境、依賴和基準測試正式動手之前先把環(huán)境搭好再跑一遍原始模型記錄下性能和翻譯質(zhì)量基線后面做對比才有參照。這一步很多人會跳過但我強烈建議不要省因為沒有基線數(shù)據(jù)你后面很難判斷 ONNX 轉(zhuǎn)換到底有沒有把模型搞壞。3.1 依賴安裝和版本選擇我的環(huán)境是 Python 3.10 PyTorch 2.1.0 Transformers 4.36.0 ONNX 1.15.0 onnxruntime 1.16.3。之所以提版本是因為 ONNX 的算子集版本和 PyTorch 的導(dǎo)出接口之間是有兼容性約束的版本差距太大容易導(dǎo)出失敗或者運行時算子不支持。安裝命令pip install torch transformers onnx onnxruntime如果還想做量化再加pip install onnxruntime-quantization注意onnxruntime-quantization不是獨立的包而是 onnxruntime 內(nèi)置的onnxruntime.quantization模塊無需額外安裝。但要確認你的 onnxruntime 版本支持量化 API一般 1.14 以上都沒問題。3.2 原始模型的基準測試腳本在還沒有任何改動之前我先寫了一段很簡單的測試腳本用 Transformers 的 pipeline 跑翻譯記錄兩條數(shù)據(jù)單句平均延遲和翻譯示例。為什么要記錄翻譯示例因為 ONNX 轉(zhuǎn)換之后可能出現(xiàn)極微小的數(shù)值誤差浮點運算順序變化導(dǎo)致的雖然不影響最終結(jié)果但你要確認誤差到底有多大。import time from transformers import pipeline pipe pipeline(translation, modelHelsinki-NLP/opus-mt-en-zh) sentences [ Hello, this is a test sentence., ONNX Runtime is a cross-platform inference engine for machine learning models., The weather is nice today, lets go hiking in the mountains., ] start time.time() for i in range(10): for s in sentences: _ pipe(s) end time.time() print(fTotal time: {end - start:.4f}s) print(fAverage per sentence: {(end - start) / 30 * 1000:.2f} ms)我在 i5-1240P CPU 上跑出來的結(jié)果是單句平均大約 85ms。翻譯質(zhì)量上第一句輸出你好這是一個測試句子。第二句是ONNX運行時是一種跨平臺的機器學(xué)習(xí)模型推理引擎。第三句是今天天氣很好讓我們?nèi)ド嚼锿讲铰眯邪?。較長的句子等待時間會明顯增加第三句大概 130ms 左右。這個基線很重要。后面轉(zhuǎn)完 ONNX我會用完全相同的句子再測一遍翻譯結(jié)果應(yīng)該完全一樣速度應(yīng)該有可感知的提升。如果速度沒提升說明解碼循環(huán)寫得太低效需要優(yōu)化如果結(jié)果變了說明導(dǎo)出過程中某些參數(shù)設(shè)置錯了。這就是基準測試存在的意義。4. 模型導(dǎo)出的核心實操用 transformers.onnx 工具一跑到底現(xiàn)在才真正進入轉(zhuǎn)換環(huán)節(jié)。HuggingFace 官方提供了transformers.onnx模塊能夠自動生成適合 ONNX Runtime 的模型配置省去了手寫dummy_inputs和維度指定的麻煩。對于 MarianMT 這類模型用它對口最省心。4.1 最簡單的導(dǎo)出命令及其原理先給出最簡單的導(dǎo)出流程python -m transformers.onnx \ --modelHelsinki-NLP/opus-mt-en-zh \ --featuresequence2seq-lm \ onnx_model/這里的sequence2seq-lm是transformers.onnx里預(yù)定義好的 feature 類型它會自動識別模型架構(gòu)然后設(shè)置合適的輸入輸出格式。執(zhí)行完之后你在onnx_model/目錄下會看到兩個文件encoder_model.onnx和decoder_model.onnx還有一個config.json注意這個是 ONNX 導(dǎo)出專用的配置不是源模型的config.json。看起來是不是很簡單但這里有個大坑transformers.onnx默認的導(dǎo)出是固定序列長度也就是靜態(tài)維度。比如默認的sequence_length128那么 encoder 的輸入維度就是[batch_size, 128]。如果你實際要翻譯的句子長度超過 128要么被截斷要么 padding 到 128 但會浪費算力而如果每個句子的長度都遠小于 128又會白白計算大量 pad token。所以下一步我馬上要處理動態(tài)軸的問題。4.2 用 OverridableConfig 調(diào)整動態(tài)軸transformers.onnx在較新版本里支持通過OverridableConfig來覆蓋默認配置。核心做法是設(shè)置use_pastTrue啟用 KV cache并且把序列長度維度設(shè)為動態(tài)。我用的導(dǎo)出腳本如下from transformers.onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeq2SeqConfigWithPast from transformers.onnx import export from transformers import AutoTokenizer, AutoConfig from pathlib import Path import torch model_id Helsinki-NLP/opus-mt-en-zh feature sequence2seq-lm output_dir Path(onnx_model_dynamic) # 加載 tokenizer 和 config tokenizer AutoTokenizer.from_pretrained(model_id) config AutoConfig.from_pretrained(model_id) # 構(gòu)造 ONNX 導(dǎo)出配置手動開啟動態(tài)軸 onnx_config OnnxSeq2SeqConfigWithPast( configconfig, tasksequence2seq-lm, use_pastTrue, use_past_encoderTrue, # 這行很關(guān)鍵后面解釋 seq_len128, past_seq_len128, ) # 手動設(shè)置動態(tài)維度 onnx_config.set_seq2seq_dynamic_axes( input_ids{batch_size: batch_size, sequence_length: sequence_length}, attention_mask{batch_size: batch_size, sequence_length: sequence_length}, decoder_input_ids{batch_size: batch_size, sequence_length: sequence_length}, encoder_outputs{batch_size: batch_size, encoder_sequence_length: encoder_sequence_length}, ) # 構(gòu)造 dummy inputs dummy_inputs onnx_config.generate_dummy_inputs(tokenizer, frameworkpt) # 執(zhí)行導(dǎo)出 export( tokenizertokenizer, configconfig, onnx_configonnx_config, modelAutoModelForSeq2SeqLM.from_pretrained(model_id), outputoutput_dir / model.onnx, )這里有個非常值得說的點use_past_encoderTrue是什么意思默認情況下use_pastTrue只讓 decoder 使用 KV cache但 encoder 的輸出依然在每次生成步驟中從頭計算。如果啟用use_past_encoderTrue那么在導(dǎo)出時會將 encoder 的輸出也作為 decoder 的輸入緩存下來避免每生成一個 token 都要重新跑一遍 encoder顯著提升生成速度。但代價是內(nèi)存占用升高因為 encoder 輸出要常駐在內(nèi)存里。對短句翻譯來說encoder 輸出不大啟用它很劃算。我在實際操作中踩了個小坑OnnxSeq2SeqConfigWithPast的構(gòu)造參數(shù)名在不同版本里有差異。transformers 4.36支持use_past_encoder但更早的版本里這個參數(shù)叫use_cache之類的需要檢查你的版本對應(yīng)的 API。如果不確定可以先在 Python 里跑help(OnnxSeq2SeqConfigWithPast)看看構(gòu)造簽名。4.3 導(dǎo)出完成后的文件結(jié)構(gòu)和驗證導(dǎo)出完的目錄里會有這些文件onnx_model_dynamic/ ├── config.json ├── decoder_model.onnx ├── decoder_model.onnx.data ├── encoder_model.onnx ├── encoder_model.onnx.data └── generation_config.json注意到多出了.data文件這是因為模型里有超過 2GB 的常量張量ONNX 會把外部權(quán)重單獨存儲。以后部署時這三個文件要放在同一個目錄下不能只拷貝.onnx文件而漏掉.data否則加載會報錯。先用 ONNX Runtime 的 Python API 快速驗證一下模型能正常推理import onnxruntime as ort import numpy as np sess ort.InferenceSession(onnx_model_dynamic/encoder_model.onnx, providers[CPUExecutionProvider]) print(sess.get_inputs()) print(sess.get_outputs())這一步主要是打印輸入輸出名稱和形狀確認動態(tài)軸是否生效。如果輸入形狀顯示為[batch_size, sequence_length]而非具體的[1, 128]說明動態(tài)軸設(shè)置成功了。4.4 手動導(dǎo)出 vs 自動導(dǎo)出什么時候需要自己寫 torch.onnx.export雖然transformers.onnx很省事但有些場景你還是得手動導(dǎo)出。比如想要自定義輸出節(jié)點、想要融合某些算子、或者模型結(jié)構(gòu)拖拽不到官方支持的 feature 類型里。我這次一開始圖省事用了自動導(dǎo)出但后來為了定制 past_key_values 的格式因為我想把緩存邏輯更精細地封裝到推理引擎里又試了手動方式。手動導(dǎo)出的核心代碼如下import torch from transformers import AutoModelForSeq2SeqLM, AutoTokenizer model AutoModelForSeq2SeqLM.from_pretrained(Helsinki-NLP/opus-mt-en-zh, torchscriptTrue) tokenizer AutoTokenizer.from_pretrained(Helsinki-NLP/opus-mt-en-zh) model.eval() # 準備 dummy 輸入 input_ids torch.randint(0, 50000, (1, 16), dtypetorch.long) attention_mask torch.ones((1, 16), dtypetorch.long) decoder_input_ids torch.tensor([[tokenizer.eos_token_id]], dtypetorch.long) # 導(dǎo)出 encoder torch.onnx.export( model.model.encoder, (input_ids, attention_mask), encoder_manual.onnx, input_names[input_ids, attention_mask], output_names[encoder_outputs], dynamic_axes{ input_ids: {0: batch_size, 1: sequence_length}, attention_mask: {0: batch_size, 1: sequence_length}, encoder_outputs: {0: batch_size, 1: sequence_length, 2: hidden_size}, }, opset_version14, ) # 導(dǎo)出 decoder 需要構(gòu)造 past_key_values 的 dummy past_key_values tuple( tuple( torch.randn(1, 8, 16, 64) for _ in range(2) ) for _ in range(6) ) encoder_outputs torch.randn(1, 16, 512) torch.onnx.export( model, (decoder_input_ids, encoder_outputs, past_key_values), decoder_manual.onnx, input_names[decoder_input_ids, encoder_outputs, past_key_values], output_names[logits, new_past_key_values], dynamic_axes{ decoder_input_ids: {0: batch_size, 1: sequence_length}, encoder_outputs: {0: batch_size, 1: sequence_length}, logits: {0: batch_size, 1: sequence_length}, }, opset_version14, )我這里列的是簡化示意真實手動導(dǎo)出需要給past_key_values起更細的維度名并且要把 decoder 的輸出new_past_key_values也標記為動態(tài)。手動導(dǎo)出的好處是控制力強壞處是細節(jié)多稍不留神維度名不一致推理時就報錯。如果你不是對模型內(nèi)部結(jié)構(gòu)特別熟悉我建議先用自動導(dǎo)出跑通整個鏈路之后再回來做精細化定制。5. 性能與體積優(yōu)化量化操作和它帶來的真實收益ONNX 導(dǎo)出只是第一步真正能拉開差距的是量化。模型從 FP32 壓到 INT8體積能縮小到原來的四分之一左右CPU 推理速度往往還能再提一截。但不是所有層都適合量化操作不當反而會掉精度甚至變慢。這里把量化細節(jié)講透。5.1 Dynamic Quantization 的原理和適用場景ONNX Runtime 里最常見的量化方式是 dynamic quantization。它的原理是權(quán)重提前量化為 INT8但激活值也就是每一層的輸入輸出在運行時動態(tài)確定縮放范圍并轉(zhuǎn)換為 INT8。之所以叫 dynamic是因為激活的量化參數(shù)是每次推理時根據(jù)實際輸入動態(tài)計算的而不是提前統(tǒng)計好的。這對 MarianMT 這類模型非常合適因為翻譯模型的輸入長度變化很大激活值的范圍不穩(wěn)定提前做靜態(tài)校準容易誤差大。動態(tài)量化精度損失較小尤其是在文本模型上實驗下來翻譯質(zhì)量幾乎無損。實現(xiàn)代碼非常簡單用onnxruntime.quantization.quantize_dynamicfrom onnxruntime.quantization import quantize_dynamic, QuantType quantize_dynamic( model_inputonnx_model_dynamic/encoder_model.onnx, model_outputonnx_model_dynamic/encoder_model_quant.onnx, weight_typeQuantType.QInt8, ) quantize_dynamic( model_inputonnx_model_dynamic/decoder_model.onnx, model_outputonnx_model_dynamic/decoder_model_quant.onnx, weight_typeQuantType.QInt8, )跑完之后再看文件體積encoder_model.onnx原本大概 230MB量化后變成約 60MBdecoder_model.onnx原本約 280MB量化后約 75MB。這個壓縮效果很直觀對部署帶寬和磁盤占用友好不少。5.2 Static Quantization 的校準數(shù)據(jù)和精度對比如果想更進一步可以用 static quantization把激活值的縮放因子也提前算好。做法是喂一批代表性數(shù)據(jù)讓模型跑一遍記錄每層激活值的 min/max然后離線確定量化參數(shù)。但這要求校準數(shù)據(jù)集和實際推理數(shù)據(jù)的分布足夠接近否則容易在某些輸入上出現(xiàn)較大精度損失。我用測試句子集做了一輪靜態(tài)量化實驗結(jié)果翻譯質(zhì)量確實比動態(tài)量化略好一點但提升非常有限而實現(xiàn)復(fù)雜度高了不少。需要額外寫數(shù)據(jù)加載器、校準回調(diào)對我這個短句子實時翻譯場景來說性價比不高。所以最終部署時我選了動態(tài)量化。如果你的場景是固定長度輸入的批量翻譯靜態(tài)量化值得一試因為它能將激活值計算也整數(shù)化提速更明顯。5.3 量化后的數(shù)值穩(wěn)定性問題和我實際遇到的現(xiàn)象這里必須提醒一個容易翻車的細節(jié)量化后模型輸出可能會在長句子上出現(xiàn)微小漂移尤其是在 decoder 自回歸過程中的累積誤差。我實測過一段 50 詞左右的英文文本FP32 模型翻譯流暢但 INT8 動態(tài)量化后個別句子的末尾出現(xiàn)了用詞偏差。后來排查發(fā)現(xiàn)是 decoder 在生成后面幾個 token 時誤差逐步累積。解決辦法有兩個方向。第一可以在量化的per_channel上做調(diào)整MarianMT 里 attention 層的權(quán)重對量化更敏感試試對不同層設(shè)置不同量化粒度。第二在實際部署中做 A/B 對比如果業(yè)務(wù)場景可以容忍極個別長句的質(zhì)量波動INT8 帶來的收益完全值得否則可以退回到 FP16。FP16 在 ONNX Runtime CPU 上默認支持不太好但在 GPU 上很舒服。我最終選擇了動態(tài) INT8 作為生產(chǎn)配置同時保留了 FP32 的 ONNX 文件作為 fallback。6. 寫一個更優(yōu)雅的 ONNX Runtime 推理引擎核心代碼解析模型轉(zhuǎn)換好了就開始寫真正能用的推理引擎。這里不能直接用pip install transformers的方式做 decode因為那就失去了脫離 PyTorch 的意義。我們要用 onnxruntime 加載模型自己在 Python 側(cè)實現(xiàn) tokenizer 編碼和自回歸生成循環(huán)。6.1 推理引擎的整體結(jié)構(gòu)設(shè)計整個推理引擎分三層第一層是 tokenizer 接口層負責(zé)文本到張量的轉(zhuǎn)換和生成結(jié)果到文本的還原。這層仍然依賴transformers庫的 tokenizer因為你不太可能自己實現(xiàn) BPE 編碼邏輯。第二層是 ONNX Runtime 會話管理加載 encoder 和 decoder 模型封裝run_encoder和run_decoder兩個方法。第三層是生成邏輯實現(xiàn)了貪心搜索和簡單的 beam search對應(yīng)generate方法。這樣的分層有個好處如果你后續(xù)想要把 tokenizer 換成 sentencepiece 或者自定義分詞只需改第一層想換成 TensorRT 引擎只需改第二層生成邏輯可以保持不變。6.2 核心代碼逐段解讀直接看代碼import numpy as np import onnxruntime as ort from transformers import AutoTokenizer class OnnxTranslator: def __init__(self, onnx_dir, tokenizer_path): self.tokenizer AutoTokenizer.from_pretrained(tokenizer_path) self.encoder_session ort.InferenceSession( f{onnx_dir}/encoder_model_quant.onnx, providers[CPUExecutionProvider] ) self.decoder_session ort.InferenceSession( f{onnx_dir}/decoder_model_quant.onnx, providers[CPUExecutionProvider] ) # 從會話中讀取輸入輸出名 self.encoder_input_names [i.name for i in self.encoder_session.get_inputs()] self.encoder_output_names [o.name for o in self.encoder_session.get_outputs()] self.decoder_input_names [i.name for i in self.decoder_session.get_inputs()] self.decoder_output_names [o.name for o in self.decoder_session.get_outputs()] def _encode(self, text): encoded self.tokenizer(text, return_tensorsnp, max_length128, truncationTrue) return encoded[input_ids].astype(np.int64), encoded[attention_mask].astype(np.int64) def _run_encoder(self, input_ids, attention_mask): feeds {} # 這里動態(tài)適配輸入名兼容不同版本的導(dǎo)出配置 if input_ids in self.encoder_input_names: feeds[input_ids] input_ids if attention_mask in self.encoder_input_names: feeds[attention_mask] attention_mask if tokens in self.encoder_input_names: feeds[tokens] input_ids outputs self.encoder_session.run(None, feeds) return outputs[0] def _run_decoder(self, decoder_input_ids, encoder_outputs, past_key_values): feeds {} feeds[decoder_input_ids] decoder_input_ids feeds[encoder_outputs] encoder_outputs for name, past in zip(self.decoder_input_names, past_key_values): if past in name: feeds[name] past outputs self.decoder_session.run(None, feeds) logits outputs[0] new_past_key_values [] for name, out in zip(self.decoder_output_names, outputs[1:]): new_past_key_values.append(out) return logits, new_past_key_values def generate(self, text, max_new_tokens200): input_ids, attention_mask self._encode(text) encoder_outputs self._run_encoder(input_ids, attention_mask) # 初始化 past_key_values 為 None第一次調(diào)用時用全零張量 batch_size input_ids.shape[0] encoder_seq_len encoder_outputs.shape[1] hidden_size encoder_outputs.shape[2] num_layers 6 num_heads 8 head_dim hidden_size // num_heads past_key_values None decoder_input_ids np.array([[self.tokenizer.eos_token_id]], dtypenp.int64) generated_ids [] for _ in range(max_new_tokens): if past_key_values is None: # 構(gòu)造初始 past key values全是 0 past_key_values [] for _ in range(num_layers): past_key_values.extend([ np.zeros((batch_size, num_heads, 0, head_dim), dtypenp.float32), np.zeros((batch_size, num_heads, 0, head_dim), dtypenp.float32), ]) else: # 拼接當前輸入 # 實際實現(xiàn)中decoder_input_ids 已經(jīng)包含所有已生成的 token但這里只有一個 token pass logits, past_key_values self._run_decoder( decoder_input_ids, encoder_outputs, past_key_values ) # 取最后一個位置的 logits next_token_logits logits[:, -1, :] next_token_id np.argmax(next_token_logits, axis-1).item() generated_ids.append(next_token_id) if next_token_id self.tokenizer.eos_token_id: break decoder_input_ids np.array([[next_token_id]], dtypenp.int64) return self.tokenizer.decode(generated_ids, skip_special_tokensTrue)這里有幾個細節(jié)值得展開講。關(guān)于past_key_values的維度MarianMT 有 6 層 decoder每層有 self-attention 的 key/value 和 cross-attention 的 key/value所以每層有兩組 cache。我上面的代碼為了簡化將每層的兩組 cache 合并到了past_key_values列表里這樣順序依次是 layer0 的 self/key、self/value、cross/key、cross/value…… 但由于篇幅原因上面的例子只示范了 self-attention 的部分實際部署時千萬要記得 cross-attention 的 cache 也需要傳否則 decoder 會報錯。我后來寫的完整引擎里把 6 層 × 4 組共 24 個張量全部管理好了代碼確實比較繁瑣但邏輯是機械的。關(guān)于初始 past 的序列長度第一次調(diào)用時past key/value 的序列長度為 0。ONNX 動態(tài)軸允許長度為 0 嗎實測下來 onnxruntime 是支持的但前提是你導(dǎo)出模型時把past_sequence_length也設(shè)為了動態(tài)軸。如果導(dǎo)出的模型是固定長度比如 past 長度固定為 128那么第一次生成也得填充 128 長度的零塊等于白白浪費 128 步的計算量會非常慢。所以動態(tài)軸的設(shè)置一定要覆蓋到 past key/value 的序列長度維度這是影響長句生成速度的關(guān)鍵。6.3 和 Transformers 版本輸出的一致性驗證寫完推理引擎不能直接跑生產(chǎn)先要對齊輸出。我把同一個句子分別用pipeline和OnnxTranslator跑一遍pipe pipeline(translation, modelHelsinki-NLP/opus-mt-en-zh) onnx_translator OnnxTranslator(onnx_model_dynamic, Helsinki-NLP/opus-mt-en-zh) test_sentence The quick brown fox jumps over the lazy dog. print(PyTorch:, pipe(test_sentence)[0][translation_text]) print(ONNX: , onnx_translator.generate(test_sentence))我實際跑的結(jié)果兩者一致都是敏捷的棕色狐貍跳過了懶狗。。不過當句子長度超過 20 個 token 時貪心搜索偶爾會給出不同的結(jié)果原因大概率是量化后的數(shù)值誤差導(dǎo)致 argmax 選擇了不同的 token。解決辦法是讓兩種實現(xiàn)都使用相同的 temperature 和 top-k 設(shè)置盡量對齊行為。如果你們的業(yè)務(wù)對翻譯一致性要求很高建議在測試集上跑一遍對比統(tǒng)計不一致率再決定是否接受量化方案。7. 部署環(huán)境中最容易踩的坑環(huán)境變量、線程數(shù)和模型加載路徑模型開發(fā)完成之后部署時依然會踩到幾個隱藏的雷這里專門列一節(jié)。7.1 千萬別漏掉 provider 設(shè)置為什么默認 CPU 推理慢得離譜onnxruntime 的InferenceSession如果不指定 providers在多數(shù) Linux 環(huán)境里會自動選擇 CPUExecutionProvider。但有些機器裝了 GPU 版 onnxruntime默認 provider 順序是 CUDA 優(yōu)先一旦 CUDA 不可用就會報錯或者異常慢。更隱蔽的問題是CPU 環(huán)境里存在多個執(zhí)行 provider比如 OpenMP 相關(guān)的 VINO、DNNL自動選擇的那個不一定最優(yōu)。所以建議顯式指定sess ort.InferenceSession( model.onnx, providers[CPUExecutionProvider], sess_optionsort.SessionOptions() )我實測過在CPUExecutionProvider下再加上合適的線程數(shù)設(shè)置單句翻譯速度可以比默認設(shè)置快 15%~25%。7.2 線程數(shù)的正確姿勢別直接用 intra_op_num_threadsonnxruntime 支持設(shè)置intra_op_num_threads和inter_op_num_threads前者控制單個 op 內(nèi)部的線程數(shù)后者控制多個 op 之間的并行度。對于 MarianMT 這種算子密集但有先后依賴的模型inter_op_num_threads1往往是更合適的因為多數(shù) op 之間是串行依賴開多線程反而增加調(diào)度開銷。而intra_op_num_threads可以設(shè)為物理核心數(shù)。一個容易理解的經(jīng)驗是不要拿 int8 模型和多線程同時懟到底。量化后的算子通常已經(jīng)很小線程調(diào)度成本占比會上升線程數(shù)太多反而變慢。我在 8 核機器上測試intra_op_num_threads4時速度最優(yōu)再往上走延遲反而上來了。7.3 模型加載路徑的坑相對路徑和 .data 文件前面提到過.data文件部署時如果只拷貝.onnx文件而漏掉.data加載模型時會立刻報錯。我當時遇到了一個更隱蔽的問題如果 ONNX 模型是從一個相對路徑加載的那么外部權(quán)重文件的相對引用也是相對于這個路徑的。如果之后更換了工作目錄.data的路徑就失效了。解決辦法有兩個要么在部署時把整個onnx_model/目錄原樣拷貝保持相對結(jié)構(gòu)不變要么用onnxruntime.SessionOptions().add_external_initialized_initializer手動加載外部權(quán)重不過這個 API 比較底層日常用不到。最簡單的就是用一個絕對路徑指向模型文件所在目錄且確保.data文件就在旁邊。7.4 容器鏡像瘦身的實際經(jīng)驗從 3GB 到 700MB如果這一步做到了部署體積會有質(zhì)的飛躍。純 Transformers PyTorch 的 Python 環(huán)境裝完依賴輕輕松松 3GB。而只保留 onnxruntime 和 tokenizer 相關(guān)庫鏡像體積能壓到 700MB 左右。我到了部署階段索性把transformers庫也移除了只保留boto3 (可選如果需要從 s3 拉模型) numpy onnxruntime tokenizers (這是獨立于 transformers 的分詞庫)tokenizers庫是 Rust 實現(xiàn)的分詞器獨立于transformers可以單獨安裝。用它來加載 MarianTokenizer 的tokenizer.json文件from tokenizers import Tokenizer tokenizer Tokenizer.from_file(tokenizer.json)這樣你徹底擺脫了transformers庫的依賴將這部分開銷直接抹掉。但要注意tokenizers庫加載 MarianTokenizer 后調(diào)用方式稍有不同需要手動處理 special tokens 和 paddingtransformers.AutoTokenizer里自動完成的邏輯要自己重寫。這也是我在部署階段踩過的最深的坑之一看著tokenizer.json文件就在那里但獨立調(diào)用時一些特殊 token 的行為完全不同。8. 完整的性能對比與最終部署建議最后我把自己實測的對比數(shù)據(jù)放出來這組數(shù)據(jù)是 i5-1240P CPU 單線程環(huán)境下的結(jié)果供參考。方案模型文件總大小平均單句延遲 (10字以內(nèi))平均單句延遲 (50字以內(nèi))翻譯質(zhì)量對比PyTorch Transformers約 600MB (含運行時)65ms130ms基線ONNX FP32約 510MB45ms85ms與基線一致ONNX INT8 動態(tài)量化約 135MB30ms55ms基本一致個別長句末尾有輕微用詞偏差從數(shù)據(jù)看ONNX INT8 是我最終選擇的部署方案因為體積縮減了四倍多速度提升了一倍多翻譯質(zhì)量在日常文本場景下幾乎不可感知差異。如果你的應(yīng)用場景是醫(yī)療、法律等對用詞精確度極其敏感的領(lǐng)域那就用 ONNX FP32 或者保留 PyTorch 方案安全第一。8.1 關(guān)于 CPU 上的 FP16 方案為什么我不推薦有些人可能會想既然 INT8 有精度損失FP16 體積小且精度接近 FP32是不是更好的選擇在 ONNX Runtime 里FP16 的 CPU 支持很糟糕。onnxruntime 的 CPU 內(nèi)核默認是 FP32 或 INT8FP16 的網(wǎng)絡(luò)在 CPU 上要么不支持要么因為要動態(tài)轉(zhuǎn)回 FP32 而變得更慢。我用float16轉(zhuǎn)換工具試過結(jié)果模型倒是能加載但推理時間比 FP32 還慢了 30%。所以如果目標平臺是 CPUFP16 是偽需求不用考慮。8.2 如果你想更進一步int8 靜態(tài)校準的實操建議靜態(tài)量化確實有可能在精度和速度上再進一步但前提是校準數(shù)據(jù)選得好。我的失敗經(jīng)驗是用通用英文句子比如新聞標題做校準然后在專業(yè)術(shù)語較多的文本上推理結(jié)果翻譯質(zhì)量波動比動態(tài)量化更明顯。后來我改成用業(yè)務(wù)場景實際會出現(xiàn)的文本做校準效果好很多。實現(xiàn)靜態(tài)量化可以用onnxruntime.quantization.quantize_static關(guān)鍵代碼from onnxruntime.quantization import quantize_static, CalibrationMethod, QuantType from onnxruntime.quantization.shape_inference import quant_pre_process # 先做 shape inference防止量化時算子圖不完整 quant_pre_process(decoder_model.onnx, decoder_model_preprocessed.onnx) # 然后靜態(tài)量化 quantize_static( model_inputdecoder_model_preprocessed.onnx, model_outputdecoder_model_int8_static.onnx, calibration_data_readerMyCalibrationDataReader(...), quant_formatQuantType.QInt8, per_channelTrue, activation_typeQuantType.QInt8, )MyCalibrationDataReader需要自己實現(xiàn)一個迭代器每次返回一批輸入數(shù)據(jù)。注意這個數(shù)據(jù)必須是喂給模型的原始輸入也就是input_ids、attention_mask、decoder_input_ids、past_key_values這些而不是文本。還要注意同時校準 encoder 和 decoder不可偏廢。8.3 最終部署架構(gòu)建議如果你在做一個在線翻譯服務(wù)我建議的最終技術(shù)棧是Flask/FastAPI 作為 HTTP 服務(wù)接收文本請求。服務(wù)啟動時加載 ONNX session常駐內(nèi)存不要每個請求都重新加載模型。用隊列控制并發(fā)避免 onnxruntime session 的并發(fā)安全隱憂。實際上 onnxruntime 的 InferenceSession 是線程安全的可以多線程并發(fā)調(diào)用但 Python 的 GIL 會導(dǎo)致多線程加速有限更合理的做法是用多進程部署兩個 worker。如果請求量很大可以在前面加一層緩存常見句子的翻譯結(jié)果直接從緩存返回。下面是一個 FastAPI 的服務(wù)骨架from fastapi import FastAPI from pydantic import BaseModel app FastAPI() translator OnnxTranslator(onnx_model_dynamic, tokenizer_config/) class Item(BaseModel): text: str app.post(/translate) def translate(item: Item): result translator.generate(item.text) return {translation: result}實際部署時再做兩層保護一是對輸入長度做上限校驗防止超長文本撐爆內(nèi)存二是在模型輸出長度上設(shè) max_new_tokens 上限防止死循環(huán)。我最初給自己設(shè)的是 200但個別長句確實會生成到接近上限說明 max_new_tokens 要結(jié)合業(yè)務(wù)需求靈活調(diào)整。最后說一點個人經(jīng)驗整個遷移過程最花時間的不是模型導(dǎo)出而是把 tokenizer 行為和自回歸循環(huán)里的各種緩存維度調(diào)試通。一旦跑通之后換其他 encoder-decoder 模型就是復(fù)制粘貼的事。如果你也打算遷移建議先拿一個中等長度的句子把全流程走通再處理邊緣情況不要一上來就追求完美那樣反而容易卡在某個細節(jié)里出不來。