1 相關預備知識
- 模型：包含了大量引數的一個網路（引數+結構），體積10MB-10GB不等
- 模型格式：相同的模型可以有不同的儲存格式（可類比音視訊檔案），目前主流有torch、tf、onnx和trt，其中tf又包含了三種格式
- 模型推理：輸入和網路中的引數進行各種運算從而得到一個輸出，計算密集型任務且需要GPU加速
- 模型推理引擎：模型推理工具，可以讓模型推理速度變快，使用該工具往往需要特定的模型格式，現在主流推理引擎有trt和ort
- 模型推理框架：對模型推理的過程進行了封裝，使之新增、刪除、替換模型更方便，更高階的框架還會有負載均衡、模型監控、自動生成grpc、http介面等功能，就是為部署而生的
接下來要介紹的triton就是目前比較優秀的一個模型推理框架。
2 從青銅到黃金：跑通triton
接下來手把手教你跑通triton，讓你明白triton到底是幹啥的。
2.1 註冊NGC平臺
NGC可以理解是NV的一個官方軟體倉庫，裡面有好多編譯好的軟體、docker映象等。我們要註冊NGC並生成相應的api key，這個api key用於在docker上登入ngc並下載裡面的映象。
註冊申請流程可以參考官方教程
2.2 登入
命令列介面輸入docker login nvcr.io
然後輸入使用者名稱和你上一步生成的key，使用者名稱就是$oauthtoken，不要忘記$符號，不要使用自己的使用者名稱。
最後會出現Login Succeeded字樣，就代表登入成功了。
2.3 拉取映象
docker pull nvcr.io/nvidia/tritonserver:22.04-py3
你也可以選擇拉取其他版本的triton。映象大概有幾個G，需耐心等待，這個映象不區分gpu和cpu，是通用的。
2.4 構建模型目錄
執行命令mkdir -p /home/triton/model_repository/fc_model_pt/1。
其中/home/triton/model_repository就是你的模型倉庫，所有的模型都在這個模型目錄中。啟動容器時會將其對映到容器中的/model資料夾上，fc_model_pt可以理解為是某一個模型的存放目錄，比如一個用於情感分類的模型，名字則沒有要求，最好見名知義，1代表版本是1
模型倉庫的目錄結構如下：
<model-repository-path>/# 模型倉庫目錄
<model-name>/ # 模型名字
[config.pbtxt] # 模型配置檔案
[<output-labels-file> ...] # 標籤檔案，可以沒有
<version>/ # 該版本下的模型
<model-definition-file>
<version>/
<model-definition-file>
...
<model-name>/
[config.pbtxt]
[<output-labels-file> ...]
<version>/
<model-definition-file>
<version>/
<model-definition-file>
...
...
2.5 生成一個用於推理的torch模型
建立一個torch模型，並使用torchscript儲存：
import torch
import torch.nn as nn
class SimpleModel(nn.Module):
def __init__(self):
super(SimpleModel, self).__init__()
self.embedding = nn.Embedding(100, 8)
self.fc = nn.Linear(8, 4)
self.fc_list = nn.Sequential(*[nn.Linear(8, 8) for _ in range(4)])
def forward(self, input_ids):
word_emb = self.embedding(input_ids)
output1 = self.fc(word_emb)
output2 = self.fc_list(word_emb)
return output1, output2
if __name__ == "__main__":
model = SimpleModel()
ipt = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long)
script_model = torch.jit.trace(model, ipt, strict=True)
torch.jit.save(script_model, "model.pt")
生成模型後拷貝到剛才建立的目錄中,注意要放到版本號對應的目錄，且模型檔名必須是model.pt。
triton支援很多模型格式，這裡只是拿了torch舉例。
2.6 編寫配置檔案
為了讓框架識別到剛放入的模型，我們需要編寫一個配置檔案config.pbtxt，具體內容如下：
name: "fc_model_pt" # 模型名，也是目錄名
platform: "pytorch_libtorch" # 模型對應的平臺，本次使用的是torch，不同格式的對應的平臺可以在官方文件找到
max_batch_size : 64 # 一次送入模型的最大bsz，防止oom
input [
{
name: "input__0" # 輸入名字，對於torch來說名字於程式碼的名字不需要對應，但必須是<name>__<index>的形式，注意是2個下劃線，寫錯就報錯
data_type: TYPE_INT64 # 型別，torch.long對應的就是int64,不同語言的tensor型別與triton型別的對應關係可以在官方文件找到
dims: [ -1 ] # -1 代表是可變維度,雖然輸入是二維的，但是預設第一個是bsz,所以只需要寫後面的維度就行(無法理解的操作，如果是[-1,-1]呼叫模型就報錯)
}
]
output [
{
name: "output__0" # 命名規範同輸入
data_type: TYPE_FP32
dims: [ -1, -1, 4 ]
},
{
name: "output__1"
data_type: TYPE_FP32
dims: [ -1, -1, 8 ]
}
]
這個模型配置檔案估計是整個triton最複雜的地方，上線模型的大部分工作估計都在寫配置檔案，我這邊三言兩語難以解釋清楚，只能給大家簡單介紹下，具體內容還請參考官方文件。注意配置檔案不要放到版本號的目錄裡，放到模型目錄裡，也就是說config.pbtxt和版本號目錄是平級的。
關於輸入shape預設第一個是bsz的官方說明：
As discussed above, Triton assumes that batching occurs along the first dimension which is not listed in in the input or output tensor dims. However, for shape tensors, batching occurs at the first shape value. For the above example, an inference request must provide inputs with the following shapes.
2.7 建立容器並啟動
執行命令：
docker run --rm -p8000:8000 -p8001:8001 -p8002:8002 \
-v /home/triton/model_repository/:/models \
nvcr.io/nvidia/tritonserver:22.04-py3 \
tritonserver \
--model-repository=/models
如果你的系統有一個可用的gpu，那你可以加上如下引數 --gpus=1來讓推理框架使用GPU加速，這個引數要放到run的後面。
2.8 測試介面
如果你按照我的教程一步一步走下去，那麼肯定可以成功啟動容器，下面我們可以寫段程式碼測試下介面是否是通的，呼叫地址是：
http:\\localhost:8000/v2/models/{model_name}/versions/{version}/infer
測試程式碼如下：
import requests
if __name__ == "__main__":
request_data = {
"inputs": [{
"name": "input__0",
"shape": [2, 3],
"datatype": "INT64",
"data": [[1, 2, 3],[4,5,6]]
}],
"outputs": [{"name": "output__0"}, {"name": "output__1"}]
}
res = requests.post(url="http://localhost:8000/v2/models/fc_model_pt/versions/1/infer",json=request_data).json()
print(res)
執行程式碼後會得到相應的輸出：
{
'model_name': 'fc_model_pt',
'model_version': '1',
'outputs': [{
'name': 'output__0',
'datatype': 'FP32',
'shape': [2, 3, 4],
'data': [1.152763843536377, 1.1349767446517944, -0.6294105648994446, 0.8846281170845032, 0.059508904814720154, -0.06066855788230896, -1.497096061706543, -1.192716121673584, 0.7339693307876587, 0.28189709782600403, 0.3425392210483551, 0.08894850313663483, 0.48277992010116577, 0.9581012725830078, 0.49371692538261414, -1.0144696235656738, -0.03292369842529297, 0.3465275764465332, -0.5444514751434326, -0.6578375697135925, 1.1234807968139648, 1.1258794069290161, -0.24797165393829346, 0.4530307352542877]
}, {
'name': 'output__1',
'datatype': 'FP32',
'shape': [2, 3, 8],
'data': [-0.28994596004486084, 0.0626179575920105, -0.018645435571670532, -0.3376324474811554, -0.35003775358200073, 0.2884367108345032, -0.2418503761291504, -0.5449661016464233, -0.48939061164855957, -0.482677698135376, -0.27752232551574707, -0.26671940088272095, -0.2171783447265625, 0.025355860590934753, -0.3266356587409973, -0.06301657110452652, -0.1746724545955658, -0.23552510142326355, 0.10773542523384094, -0.4255935847759247, -0.47757795453071594, 0.4816707670688629, -0.16988427937030792, -0.35756853222846985, -0.06549499928951263, -0.04733048379421234, -0.035484105348587036, -0.4210450053215027, -0.07763291895389557, 0.2223128080368042, -0.23027443885803223, -0.4195460081100464, -0.21789231896400452, -0.19235755503177643, -0.16810789704322815, -0.34017443656921387, -0.05121977627277374, 0.08258339017629623, -0.2775516211986542, -0.27720844745635986, -0.25765007734298706, -0.014576494693756104, 0.0661710798740387, -0.38623639941215515, -0.45036202669143677, 0.3960753381252289, -0.20757021009922028, -0.511818528175354]
}]
}
不知道是不是我的用法有問題，就使用體驗來看，這個推理介面有些讓我不適應：
- 明明在config.pbtxt裡指定了datatype,但是輸入的時候還需要指定，不指定就會報錯
- 輸入的shape也需要指定，那不然也會報錯
- datatype的值和config.pbtxt裡不統一，如果datatype設為TYPE_INT64，則會報錯，必須為INT64
- 輸出裡的data是個1維陣列，需要根據返回的shape自動reshape成對應的array
除了像我這樣直接寫程式碼呼叫，還可以使用他們提供的官方庫pip install tritonclient[http],地址如下：https://github.com/triton-inference-server/client。
3 從黃金到王者：使用triton的高階特性
上一小節的教程只是用到了triton的基本功能，所以段位只能說是個黃金，下面介紹下一些triton的高階特性。
3.1 模型並行
模型並行可以指同時啟動多個模型或單個模型的多個例項。實現起來並不複雜，通過修改配置引數就可以。在預設情況下，triton會在每個可用的gpu上都部署一個該模型的例項從而實現並行。
接下來我會對多種情況進行測試，以讓你們清楚模型並行所帶來的效果，本人的配置是2塊3060(租的)用於測試多模型的情況。
壓測命令使用ab -k -c 5 -n 500 -p ipt.json http://localhost:8000/v2/models/fc_model_pt/versions/1/infer
這條命令的意思是5個程式反覆呼叫介面共500次。
測試配置及對應的QPS如下：
- 共1個卡；每個卡執行1個例項：QPS為603
- 共2個卡；每個卡執行1個例項：QPS為1115
- 共2個卡；每個卡執行2個例項：QPS為1453
- 共2個卡；每個卡執行2個例項；同時在CPU上放2個例項：QPS為972
結論如下：多卡效能有提升；多個例項能進一步提升併發能力；加入CPU會拖累速度，主要是因為CPU速度太慢。
下面是上述測試對應的配置項，直接複製了放到config.pbtxt中就行
#共2個卡；每個卡執行2個例項
instance_group [
{
count: 2
kind: KIND_GPU
gpus: [ 0 ]
},
{
count: 2
kind: KIND_GPU
gpus: [ 1 ]
}
]
# 共2個卡；每個卡執行2個例項；同時在CPU上放2個例項
instance_group [
{
count: 2
kind: KIND_GPU
gpus: [ 0 ]
},
{
count: 2
kind: KIND_GPU
gpus: [ 1 ]
},
{
count: 2
kind: KIND_CPU
}
]
至於選擇使用幾張卡，則通過建立容器時的--gpus來指定
3.2 動態batch
動態batch的意思是指, 對於一個請求,先不進行推理,等個幾毫秒，把這幾毫秒的所有請求拼接成一個batch進行推理，這樣可以充分利用硬體，提升並行能力，當然缺點就是個別使用者等待時間變長，不適合低頻次請求的場景。使用動態batch很簡單，只需要在config.pbtxt加上 dynamic_batching { }，具體引數細節大家可以去看文件，我的這種簡單寫法，組成的bsz上限就是max_batch_size，本人壓測的結果是約有50%QPS提升，反正就是有效果就對了。
PS：這種優化方式對於壓測來說簡直就是作弊。。。
3.3 自定義backend
所謂自定義backend就是自己寫推理過程，正常情況下整個推理過程是通過模型直接解決的，但是有一些推理過程還會包含一些業務邏輯，比如：整個推理過程需要2個模型，其中要對第一個模型的輸出結果做一些邏輯判斷然後修改輸出才能作為第二個模型的輸入，最簡單的做法就是我們呼叫兩次triton服務，先呼叫第一個模型獲取輸出然後進行業務邏輯判斷和修改，然後再呼叫第二個模型。不過在triton中我們可以自定義一個backend把整個呼叫過程寫在裡面，這樣就簡化呼叫過程，同時也避免了一部分http傳輸時延。
我舉的例子其實是一個包含了業務邏輯的pipline，這種做法NV叫做BLS(Business Logic Scripting)。
要實現自定義的backend也很簡單，與上文講的放torch模型流程基本一樣，首先建立模型資料夾，然後在資料夾裡新建config.pbtxt,然後新建版本資料夾，然後放入model.py，這個py檔案裡就寫了推理過程。為了說明目錄結構，我把構建好的模型倉庫目錄樹列印出來展示一下：
model_repository/ # 模型倉庫
|-- custom_model # 我們的自定義backend模型目錄
| |-- 1 # 版本
| | |-- model.py # 模型Py檔案，裡面主要是推理的邏輯
| `-- config.pbtxt # 配置檔案
`-- fc_model_pt # 上一小節介紹的模型
|-- 1
| `-- model.pt
`-- config.pbtxt
如果上一小節你看明白了，那麼你就會發現自定義backend的模型目錄設定和正常目錄設定是一樣的，唯一不同的就是模型檔案由網路權重變成了自己寫的程式碼而已。下面說下config.pbtxt和model.py的檔案內容，大家可以直接複製貼上：
# model.py
import json
import numpy as np
import triton_python_backend_utils as pb_utils
class TritonPythonModel:
"""Your Python model must use the same class name. Every Python model
that is created must have "TritonPythonModel" as the class name.
"""
def initialize(self, args):
"""`initialize` is called only once when the model is being loaded.
Implementing `initialize` function is optional. This function allows
the model to intialize any state associated with this model.
Parameters
----------
args : dict
Both keys and values are strings. The dictionary keys and values are:
* model_config: A JSON string containing the model configuration
* model_instance_kind: A string containing model instance kind
* model_instance_device_id: A string containing model instance device ID
* model_repository: Model repository path
* model_version: Model version
* model_name: Model name
"""
# You must parse model_config. JSON string is not parsed here
self.model_config = model_config = json.loads(args['model_config'])
# Get output__0 configuration
output0_config = pb_utils.get_output_config_by_name(
model_config, "output__0")
# Get output__1 configuration
output1_config = pb_utils.get_output_config_by_name(
model_config, "output__1")
# Convert Triton types to numpy types
self.output0_dtype = pb_utils.triton_string_to_numpy(output0_config['data_type'])
self.output1_dtype = pb_utils.triton_string_to_numpy(output1_config['data_type'])
def execute(self, requests):
"""
requests : list
A list of pb_utils.InferenceRequest
Returns
-------
list
A list of pb_utils.InferenceResponse. The length of this list must
be the same as `requests`
"""
output0_dtype = self.output0_dtype
output1_dtype = self.output1_dtype
responses = []
# Every Python backend must iterate over everyone of the requests
# and create a pb_utils.InferenceResponse for each of them.
for request in requests:
# 獲取請求資料
in_0 = pb_utils.get_input_tensor_by_name(request, "input__0")
# 第一個輸出結果自己隨便造一個假的，就假裝是有邏輯了
out_0 = np.array([1, 2, 3, 4, 5, 6, 7, 8]) # 為演示方便先寫死
out_tensor_0 = pb_utils.Tensor("output__0", out_0.astype(output0_dtype))
# 第二個輸出結果呼叫fc_model_pt獲取結果
inference_request = pb_utils.InferenceRequest(
model_name='fc_model_pt',
requested_output_names=['output__0', 'output__1'],
inputs=[in_0])
inference_response = inference_request.exec()
out_tensor_1 = pb_utils.get_output_tensor_by_name(inference_response, 'output__1')
inference_response = pb_utils.InferenceResponse(output_tensors=[out_tensor_0, out_tensor_1])
responses.append(inference_response)
return responses
def finalize(self):
"""`finalize` is called only once when the model is being unloaded.
Implementing `finalize` function is OPTIONAL. This function allows
the model to perform any necessary clean ups before exit.
"""
print('Cleaning up...')
config.pbtxt的檔案內容：
name: "custom_model"
backend: "python"
input [
{
name: "input__0"
data_type: TYPE_INT64
dims: [ -1, -1 ]
}
]
output [
{
name: "output__0"
data_type: TYPE_FP32
dims: [ -1, -1, 4 ]
},
{
name: "output__1"
data_type: TYPE_FP32
dims: [ -1, -1, 8 ]
}
]
待上述工作都完成後，就可以啟動程式檢視執行結果了，大家可以直接複製我的程式碼進行測試：
import requests
if __name__ == "__main__":
request_data = {
"inputs": [{
"name": "input__0",
"shape": [1, 2],
"datatype": "INT64",
"data": [[1, 2]]
}],
"outputs": [{"name": "output__0"}, {"name": "output__1"}]
}
res = requests.post(url="http://localhost:8000/v2/models/fc_model_pt/versions/1/infer",json=request_data).json()
print(res)
res = requests.post(url="http://localhost:8000/v2/models/custom_model/versions/1/infer",json=request_data).json()
print(res)
執行結果可以發現2次的輸出中，ouput__0是不一樣，而output__1是一樣的，這和我們們寫的model.py的邏輯有關係，我這裡就不多解釋了。
PS：自定義backend避免了需要反覆呼叫NLG模型進行生成而產生的傳輸時延
4 總結
如果上面的教程你完整的走了一遍，相信你對triton的使用方法和相關特性會有一個大概的瞭解，王者不敢保證，黃金段位肯定是有了，後面去學習使用triton的其他特性想必也會非常順利.
如果你在使用過程中遇到了問題可以私信或者評論，我們一起學習交流。
最後如果有需要租用GPU機器的同學，可以考慮featurize，我做的實驗就是在這上面租了機器，大家如果需要租用請使用我的邀請連結，也算是給我做實驗回本了，謝謝各位，註冊連結：https://featurize.cn?s=7b59a59ea4574318b0504dff01728f95
文章同步發與知乎和公眾號，歡迎關注：
https://www.zhihu.com/people/zdd-44-59
