一、背景
對於算法工程師來說,通常採用python語言來作爲工作語言,但是直接用python部署線上服務性能很差。這個問題困擾了我很久,爲了緩解深度學習模型工程落地性能問題,探索了Nvidia提供的triton部署框架,並在九數中臺上完成線上部署,發現性能提升近337%!!(原服務單次訪問模型推理時間175ms左右,同模型框架單次訪問推理時間縮減至40ms左右)。本文僅以本次部署服務經驗分享給諸多受制於python性能問題的夥伴,希望有所幫助。
二、問題發現
組內原算法服務一直採用python-backend的鏡像部署方式,將算法模型包裝成接口形式,再通過Flask外露,打入docker中啓動服務,但是發現推到線上接口響應時間過長,非常影響用戶體驗,於是想做出改進。python後端部署一般存在以下問題:
爲了解決上述困擾,我們經過調研行業內其他實踐經驗,決定摒棄傳統鏡像部署方式,擺脫python部署環境,換用triton部署算法模型。
三、部署實例
本試驗通過fine-tune的Bert模型進行文本分類任務,並通過九數算法中臺包裝的triton部署框架部署至公網可訪問,以下介紹了我們團隊部署的全流程和途中踩到的坑,以供大家快速入手(以下均以Bert模型爲例):
step 1: 將訓練好的模型保存爲onnx或torchscript(即.pt)文件格式。據調研行業經驗,onnx一般用於cpu環境下推理任務效果較好,.pt文件在gpu環境下推理效果較好,我們最終選擇了將模型轉化爲.pt文件。代碼例子如下:
import torch
from transformers import BertTokenizer, BertForSequenceClassification
tokenizer = BertTokenizer.from_pretrained('./pretrained-weights/bert-base-chinese')
model = BertForSequenceClassification.from_pretrained('./pretrained-weights/bert-base-chinese', num_labels=10)
# Trace the wrapped model
# input_ids和attention_mask來自tokenizer, 具體可看一下torch.jit.trace用法
traced_model = torch.jit.trace(model, (input_ids, attention_mask))
traced_model.save("./saved_model_torchscript_v3.pt")
本段代碼借用了torch.jit.trace的存儲方式,輸入是文本經過tokenizer之後的input_ids和attention_mask,輸出是未經softmax處理的線性層。我們嘗試將softmax部分包裝進模型中一起打包成pt文件,使模型能夠直接吐出分類類別,方案是外層再包裝一層forward,代碼例子如下:
class BertSequenceClassifierWrapper(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model.cuda()
def forward(self, input_ids, attention_mask):
outputs = self.model(input_ids=input_ids.cuda(), attention_mask=attention_mask.cuda())
# Apply softmax to the logits
probs = softmax(outputs.logits, dim=-1)
# Get the predicted class
predicted_class = torch.argmax(probs, dim=-1)
return predicted_class
wrapped_model = BertSequenceClassifierWrapper(model)
traced_model = torch.jit.trace(wrapped_model, (input_ids, attention_mask))
traced_model.save("./saved_model_torchscript_v3.pt")
然而,參考於行業內實踐,我們沒有選擇將tokenizer部分包裝至模型內部,這是因爲torchscript僅接受數組類型輸入而非文本。最終,我們封裝的模型輸入爲input_ids和attention_mask,輸出爲分類結果(標量輸出)。
【注】:坑點:
step 2: 將.pt文件存放在九數磁盤目錄下指定位置,注意,目錄格式需要嚴格按照要求,具體可參考九數幫助文檔。
目錄格式如下(其中model.pt即爲我們訓練好的模型):
pytorch-backend-online
|-- pytorch-model
| |-- 1
| | |-- model.pt
| `-- config.pbtxt
config.pbtxt配置如下:
name: "pytorch-model"
platform: "pytorch_libtorch"
input [
{
name: "INPUT__0"
data_type: TYPE_INT64
dims: [1, 512]
},
{
name: "INPUT__1"
data_type: TYPE_INT64
dims: [1, 512]
}
]
output [
{
name: "OUTPUT__0"
data_type: TYPE_INT64
dims: [1, 1]
}
]
instance_group [
{
count: 1
}
]
配置中,我們需要指定輸入和輸出的維度以及數據類型。注意,數據類型指的是每個可迭代對象的類型(比如512維向量每個維度都是int64)。在本實例中,我們定義了兩個輸入均爲1*512維向量,輸出爲一個標量。
step 3: 模型註冊。需要到九數上註冊你的模型,然後才能部署,方式如下:
模型註冊表填寫如下:
坑點:圖中step2定義的一級目錄,比如 /home/{erp}/xxx/xxx/xxx。
step 4: 模型註冊好之後,就可以在註冊好的模型下點擊:部署-測試
然後需要填寫UI界面的信息。
到這裏,如果模型沒有問題,部署就是成功了~
step 5: 測試端口通不通。如果是測試環境,我們可以通過post接口在notebook中測試接口通不通,測試demo如下:
URL來自獲取方式如下:
如果測試沒有問題,就可以轉生產了~配置好域名和端口就可以公網訪問了。至此,triton框架部署torchscript方式完結。