我不會用 Triton 系列：Triton 搭建 ensemble 過程記錄

Triton 搭建 ensemble 過程記錄

本文記錄 Triton ensemble 搭建的過程，在 Triton 這個特性叫做 ensemble，但是這個特性叫做 pipeline 更為常見，后面就叫 pipeline 吧。首先要說明的是，本文中的例子只是為了試試看 Triton pipeline 這個特性，我認為搭建出的 pipeline 不一定就是高效的。

先來說說本文將要搭建什么樣的 pipeline。本文將使用 resnet50 來進行圖片分類，分類的類別保持不變。在對圖片進行分類之前，一般都需要有一個預處理的過程。因此，這篇文章將搭建的 pipeline 很簡單，就先進行預處理，然后分類。預處理采用 DALI 來處理，resnet50 使用 Pytorch 導出。

本文使用的模型配置文件，已經放在了 Github 上面：https://github.com/zzk0/triton

Pytorch 搭建 resnet50

Pytorch 導出 resnet50 模型

非常簡單的一個代碼片段，於是我們得到了一個 torchscript 模型。

import torch
import torchvision.models as models

resnet50 = models.resnet50(pretrained=True)
resnet50.eval()
image = torch.randn(1, 3, 244, 244)
resnet50_traced = torch.jit.trace(resnet50, image)
resnet50(image)
resnet50_traced.save('model.pt')

Pytorch 模型配置

我們將文件按照如下方式進行組織，其中 config.pbtxt 是模型配置文件，labels.txt 是 resnet50 訓練時候的分類類別，里面有一千個類。另外還需要注意的是，labels.txt 里面的寫法，就是一個字符串一個類別就好了。

.
├── 1
│   └── model.pt
├── config.pbtxt
├── labels.txt

config.pbtxt 的寫法，通過指定 label_filename 來設定標簽文件，輸出有 1000 維。

name: "resnet50_pytorch"
platform: "pytorch_libtorch"
max_batch_size: 128
input [
  {
    name: "INPUT__0"
    data_type: TYPE_FP32
    dims: [ 3, -1, -1 ]
  }
]
output [
  {
    name: "OUTPUT__0"
    data_type: TYPE_FP32
    dims: [ 1000 ]
    label_filename: "labels.txt"
  }
]
instance_group [
  {
    count: 1
    kind: KIND_GPU
  }
]

客戶端

將模型放到 Triton 的模型倉庫之后，啟動服務器。之后我們使用下面的腳本進行請求。在這個客戶端里，我們先自己做預處理，后續我們將會把預處理的操作放置到服務端。

如果我們想要獲取分類的結果，我們可以設置 class_count=k，表示獲取 TopK 分類預測結果。如果沒有設置這個選項，那么將會得到一個 1000 維的向量。

import numpy as np
import tritonclient.http as httpclient
import torch
from PIL import Image


if __name__ == '__main__':
    triton_client = httpclient.InferenceServerClient(url='172.17.0.2:8000')

    image = Image.open('../resources/images/cat.jpg')
    
    image = image.resize((224, 224), Image.ANTIALIAS)
    image = np.asarray(image)
    image = image / 255
    image = np.expand_dims(image, axis=0)
    image = np.transpose(image, axes=[0, 3, 1, 2])
    image = image.astype(np.float32)

    inputs = []
    inputs.append(httpclient.InferInput('INPUT__0', image.shape, "FP32"))
    inputs[0].set_data_from_numpy(image, binary_data=False)
    outputs = []
    outputs.append(httpclient.InferRequestedOutput('OUTPUT__0', binary_data=False, class_count=1))

    results = triton_client.infer('resnet50_pytorch', inputs=inputs, outputs=outputs)
    output_data0 = results.as_numpy('OUTPUT__0')
    print(output_data0.shape)
    print(output_data0)

DALI

接下來，我們將客戶端預處理的操作放到了服務端上。這里必須要指出的是，這么做只是為了搭建 pipeline，並不是為了性能。你想，圖片沒有預處理之前，是不是很大，通過網絡傳輸到服務端的開銷可能蓋過了服務端預處理的收益。

導出 DALI 預處理 pipeline

通過下面的腳序列化 pipeline。

import nvidia.dali as dali
import nvidia.dali.fn as fn

@dali.pipeline_def(batch_size=128, num_threads=4, device_id=0)
def pipeline():
    images = fn.external_source(device='cpu', name='DALI_INPUT_0')
    images = fn.resize(images, resize_x=224, resize_y=224)
    images = fn.transpose(images, perm=[2, 0, 1])
    images = images / 255
    return images


pipeline().serialize(filename='./1/model.dali')

DALI 模型配置

我們將文件按照如下方式組織。

.
├── 1
│   └── model.dali
├── config.pbtxt

模型配置如下。需要注意一個問題：模型實例化的時候，如果沒有設置設備，Triton 會在每個設備上初始化一個，接着會發生 core dump。目前猜想的原因是，序列化保存的 pipeline 保存了 device_id=0 這個信息，然后在我的服務器上的第二張卡上初始化模型實例的時候，會出錯。后續仔細分析看看，提個 issue 或 pr。

配置好之后，放到模型倉庫，然后使用 Github 中對應的腳本做請求試試看，這里就不啰嗦了。

name: "resnet50_dali"
backend: "dali"
max_batch_size: 128
input [
  {
    name: "DALI_INPUT_0"
    data_type: TYPE_FP32
    dims: [ -1, -1, 3 ]
  }
]

output [
  {
    name: "DALI_OUTPUT_0"
    data_type: TYPE_FP32
    dims: [ 3, 224, 224 ]
  }
]
instance_group [
  {
    count: 1
    kind: KIND_GPU
    gpus: [ 0 ]
  }
]

搭建 pipeline

模型配置

pipeline 的配置方法也挺簡單的，只不過個人覺得手寫 protobuf 不太順手，用戶體驗不太好。

下面說幾個要注意的點：一，ensemble 的 key 是 platform，不是 backend。二，model_version 設為數字，而不是字符串。三，ensemble_scheduling 的輸入輸出 key 都是對應模型的輸入輸出名字。

name: "resnet50_ensemble"
platform: "ensemble"
max_batch_size: 128
input [
  {
    name: "ENSEMBLE_INPUT_0"
    data_type: TYPE_FP32
    dims: [ -1, -1, 3 ]
  }
]
output [
  {
    name: "ENSEMBLE_OUTPUT_0"
    data_type: TYPE_FP32
    dims: [ 1000 ]
  }
]
ensemble_scheduling {
  step [
    {
      model_name: "resnet50_dali"
      model_version: 1
      input_map: {
        key: "DALI_INPUT_0"
        value: "ENSEMBLE_INPUT_0"
      }
      output_map: {
        key: "DALI_OUTPUT_0"
        value: "preprocessed_image"
      }
    },
    {
      model_name: "resnet50_pytorch"
      model_version: 1
      input_map: {
        key: "INPUT__0"
        value: "preprocessed_image"
      }
      output_map: {
        key: "OUTPUT__0"
        value: "ENSEMBLE_OUTPUT_0"
      }
    }
  ]
}

客戶端請求

雖然在客戶端避開預處理，但是不能完全避開。比如我們一定需要設置好輸入的 shape，否則 Triton 就是不認你這個請求，所以還是自己手動加一個維度。此外，輸入要設置成 float32 類型。於是我們避開了 resize 等預處理操作。

請求的時候，你會發現，即使 pipeline 沒有設置 label_filename，我們仍然可以獲取分類的結果。這里我猜測 Triton 的內部實現可能是，輸入的 Shape 會進行檢查，輸出的 Shape 就不理了（這個不是看 Backend 是否檢查嘛。

import numpy as np
import tritonclient.http as httpclient
import torch
from PIL import Image


if __name__ == '__main__':
    triton_client = httpclient.InferenceServerClient(url='172.17.0.2:8000')

    image = Image.open('../resources/images/cat.jpg')
    image = np.asarray(image)
    image = np.expand_dims(image, axis=0)
    image = image.astype(np.float32)

    inputs = []
    inputs.append(httpclient.InferInput('ENSEMBLE_INPUT_0', image.shape, "FP32"))
    inputs[0].set_data_from_numpy(image, binary_data=False)
    outputs = []
    outputs.append(httpclient.InferRequestedOutput('ENSEMBLE_OUTPUT_0', binary_data=False, class_count=1))

    results = triton_client.infer('resnet50_ensemble', inputs=inputs, outputs=outputs)
    output_data0 = results.as_numpy('ENSEMBLE_OUTPUT_0')
    print(output_data0.shape)
    print(output_data0)

至此，我們就可以使用一張沒有預處理過的照片，然后直接發送給 Triton，Triton 幫你做預處理，然后d對處理的結果做分類。不過，我現在對 pipeline 的原理還不是很清楚，比如有個問題：pipeline 的模型和模型之間，是否會發生額外的內存復制開銷呢？這個要深入源碼看一看了。

附上自己的請求結果。