一、背景
對於算法工程師來說，通常採用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.pbtxtconfig.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方式完結。