TinyBERT和ONNX runtime使用心得

BERT的CPU inference 性能真是悲劇。之前做研究發現用 fine-tuned BERT在報告匯入 Elasticsearch 前preprocessing，可解決否定語句（No evidence of XXX）干擾搜尋結果的問題。研究時 BERT 很好用，準確度很高。但佈署到生產環境才發現有嚴重性能問題。

訓練時，雲端GPU租賃有RTX 3090、4090吃到飽。訓練樣本量（幾千份）很小，沒發現問題。做staging（no CUDA）和推到生產環境（Intel Core I3-9100，Intel 內顯）問題就來了。

在 staging 機器上，加了DistilBERT模型做 preprocessing 之後，文本索引速度暴跌到一小時幾千份。而目前要索引的報告量接近800萬份。生產環境基本不能用。

後來找網路攻略。發現三個解決思路：換輕量化模型、換Runtime、模型瘦身（做 INT8 quantization或pruning等，「據說」對CPU架構會友善點）。

輕量化模型找到了：華為的 TinyBERT。原始論文號稱推理速度比BERT快接近10倍。實際把DistilBERT換成TinyBERT後，索引速度在staging和生產環境就暴增到一小時2萬份，F1 score （0.94）和原本用DistilBERT差不多。但雖然性能有很大提升，一次完整索引迭代還是要超過2個禮拜。做到一半發現有問題需要重新索引又是另一個迭代。

後來繼續查，照著網路攻略把INT8 quantization弄了出來，但staging的時候跟我說沒GPU不能開quantization！功力不夠找不到問題就先放棄了。再試了換Runtime。Runtime換到ONNX以後。Jupyter notebook上推理的 wall time 直接變1/3（68000+ 測資 14min -> 5min）。於是再熱火朝天改code佈署。結果性能提升比三倍更多（原因不清楚，權當神秘學XD），索引速度爆增到每小時15萬份，一次索引迭代只要兩三天。已可接受了。

結論而言。如果未來服務要更進一步，還是要想辦法搞 GPU。用CPU跑BERT根本自虐，還好這次有找到解法，不然就得關feature然後說研究成果落地有困難QQ

全文連結