1 背景

轉轉面向二手電商業(yè)務，在搜索推薦、智能質檢、智能客服等場景落地了AI技術。在實踐的過程中，也發(fā)現了存在GPU執(zhí)行優(yōu)化不充分，浪費計算資源，增加應用成本等問題。

此外還存在線上線下處理邏輯需要分別開發(fā)的情況，造成額外的開發(fā)成本和錯誤排查成本，對一些需要高速迭代的業(yè)務場景的負面影響不可忽視。本文將會重點介紹基于Torchserve進行推理服務部署架構優(yōu)化的工程實踐，希望對面臨類似問題的同學們有所幫助。

2 問題和解決思路

2.1 現狀分析

圖片

上圖為之前的推理系統(tǒng)架構圖，采用CPU和GPU分離的架構。這種架構的特點是:

GPU部分和CPU部分解耦分別部署微服務，預處理部分一般是在CPU上執(zhí)行，容易成為推理服務的性能瓶頸。解耦后可以將CPU部分部署在單獨的機器上，無限水平擴容擴容。

圖片

上圖為美團通用高效的推理服務部署架構方案，可以看到架構思路基本相同。

2.2 問題

該方案具有很大的優(yōu)點，也是業(yè)界多個公司采取的方案，但是架構的選項也需要考慮具體的業(yè)務場景。該方案在轉轉的的場景里就出現了一些問題，比如：

• 迭代效率：GPU執(zhí)行部分的是基于torch或者tf的模型，可以接近0成本的部署到推理微服務上，但是CPU執(zhí)行部分包含一些預處理、后處理的部分，除了一些常見的圖像解碼、縮放、NMS等操作，還有很多自定義的邏輯，難以通過統(tǒng)一協(xié)議進行低成本部署，需要在微服務里進行二次開發(fā)，并且大部分情況下需要采用和離線算法實現不同的開發(fā)語言，拖累業(yè)務迭代效率。當前業(yè)界算法工程師的主要工作語言為python，對其他的語言或者dsl都存在一定的學習成本。
• 網絡通信：轉轉的部分業(yè)務場景存在圖片尺寸較大的情況，因為二手的某些質檢場景需要高清圖片判斷有沒有破損、劃痕、污漬等細小痕跡。這樣微服務之間的網絡通信開銷負擔比一般場景大很多。

2.3 解決思路

2.3.1框架調研

深度學習模型的常見的部署框架有以下幾個：

從性能和質量的角度，三個框架水平都可以達到要求。重點考慮三個框架對于非深度學習部分的自定義邏輯支持：

• tensoflow的tf.function裝飾器和AutoGraph機制 tf.function是TensorFlow 2.x中引入的一個重要特性，它通過在Python函數上應用一個裝飾器，將原生Python代碼轉換為TensorFlow圖代碼，從而享受圖執(zhí)行帶來的性能優(yōu)勢。AutoGraph是tf.function底層的一項關鍵技術，它可以將復雜的Python代碼，比如包含while，for，if的控制流，轉換為TensorFlow的圖,例如：

@tf.function
def fizzbuzz(n):
  for i in tf.range(n):
    if i % 3 == 0:
      tf.print('Fizz')
    elif i % 5 == 0:
      tf.print('Buzz')
    else:
      tf.print(i)

fizzbuzz(tf.constant(15))

• triton Python Backend機制 triton允許使用Python編寫后端邏輯，這樣可以利用Python的靈活性和豐富的庫。Python Backend通過實現initialize、execute和finalize等接口來完成模型的加載、推理和卸載。這種方式適合于需要復雜邏輯處理的場景，比如多模型協(xié)同工作或者需要自定義預處理和后處理的情況。Python Backend可以與Triton的pipeline功能結合，實現更復雜的推理流程。例如下面的代碼中，實現了一個initialize方法初始化，并且實現一個execute方法執(zhí)行具體的邏輯，代碼為python實現。

import numpy as np

class PythonAddModel:
    def initialize(self, args):
        self.model_config = args['model_config']

    def execute(self, requests):
        responses = []
        for request in requests:
            out_0 = request.inputs[0].as_numpy() + request.inputs[1].as_numpy()
            out_tensor_0 = pb_utils.Tensor("OUT0", out_0.astype(np.float32))
            responses.append(pb_utils.InferenceResponse([out_tensor_0]))
        return responses

• torchserve Custom handlers TorchServe通過定義handler來處理模型的加載、預處理、推理和后處理。handler通常繼承自BaseHandler類，并重寫initialize、preprocess、inference和postprocess等方法。如下面代碼所示，與Triton Python Backend有些類似。

from ts.torch_handler import TorchHandler

class ImageClassifierHandler(TorchHandler):
    def initialize(self, params):
        """初始化模型"""
        self.model = SimpleCNN()
        self.model.load_state_dict(torch.load('model.pth', map_location=torch.device('cuda:0')))
        self.model.eval()

    def preprocess(self, batch):
        """預處理輸入數據"""
        images = [img.convert('RGB') for img in batch]
        images = [img.resize((224, 224)) for img in images]
        images = [torch.tensor(np.array(img)).permute(2, 0, 1).float() for img in images]
        images = [img / 255.0 for img in images]
        return images

    def postprocess(self, outputs):
        """后處理輸出結果"""
        _, predicted = torch.max(outputs, 1)
        return predicted

2.3.2框架選型

tensorflow serving的tf.function裝飾器和AutoGraph機制并不能保證兼容所有的python代碼和控制流，并不滿足需求，在兼容第三方python包上也存在問題。此外tensorflow作為早年應用最廣的深度學習框架，近年來在流行度上已經有被后來追上的趨勢，tensorflow serving基本上只支持tensorflow框架，所以第一個排除。 

Triton Python Backend和torchserve Custom handlers在功能和機制上比較類似。都提供了一個靈活、易用且可擴展的解決方案，特別適合于需要快速部署和靈活處理模型的場景?？蚣芗嫒萆蟭riton支持主流的所有框架，torchserve主要支持pytorch和onnx協(xié)議，都可以滿足轉轉的需求。經過調研和試用，我們最終選擇了torchserve作為本次的框架選項，原因如下：

• torchserve與PyTorch生態(tài)深度集成，而Triton的學習曲線相對陡峭。torchserve主要支持torch框架，同時只兼容onnx協(xié)議。在輕量級和易用性上更符合轉轉當前的業(yè)務場景要求。例如在模型部署的格式轉換和配置上，torchserve相較于triton要簡易很多。
• 從后續(xù)轉轉GPU推理服務的演進來看，長期來看支持所有主流推理框架是必需的，短期在業(yè)務高速成長期優(yōu)先選擇一個框架與長期支持所有主流并不沖突。

3 torchserve實踐過程

3.1 torchserve使用與調優(yōu)

3.1.1 使用流程

圖片

以一個圖像模型簡單舉例，如圖所示：

• 將模型權重文件及前后處理邏輯python代碼打包成一個mar包
• mar包提交到torchserve進程中進行模型注冊
• 請求到來后執(zhí)行圖片下載和模型前處理、推理、后處理，返回結果

mar包打包指令：

torch-model-archiver --model-name your_model_name --version 1.0 --serialized-file path_to_your_model.pth --handler custom_handler.py --extra-files path_to_any_extra_files

• your_model_name：你為模型指定的名稱。
• 1.0：模型的版本號，可以根據實際情況進行修改。
• path_to_your_model.pth：你的 PyTorch 模型文件的路徑。
• custom_handler.py：處理模型輸入和輸出的自定義處理函數文件。
• path_to_any_extra_files：如果有其他需要一起打包的文件，可以在這里指定路徑，可以是多個文件路徑用逗號分隔。

torchserve的custom handler機制和易用性對于開發(fā)效率的提升是顯著的，在我們的內部場景里，一個單人維護的推理服務，在半年內節(jié)省了約32PD（人日）的開發(fā)成本。

3.1.2 torch-trt

模型的主干網絡部分，需要進行優(yōu)化，否則執(zhí)行效率差耗時較長。torch-trt允許將PyTorch 模型轉換為 TensorRT 格式，從而能夠利用 TensorRT 強大的優(yōu)化引擎。TensorRT 針對 NVIDIA GPU 進行了高度優(yōu)化，能夠實現快速的推理性能。它通過對模型進行層融合、內核自動調整和內存優(yōu)化等操作，顯著提高了模型的推理速度。

import torch
   import torch_tensorrt

   # Load your PyTorch model
   model = torch.load('path_to_your_model.pth')

   # Convert the model to TensorRT
   trt_model = torch_tensorrt.compile(model, inputs=[torch_tensorrt.Input((1, 3, 224, 224))], enabled_precisinotallow={torch.float32})

   # Save the converted model
   torch.save(trt_model, 'path_to_trt_model.pth')

torch-trt比起tensorflow-trt和triton-trt相對來說比較簡單：

• 使用torch.compile可以一鍵轉換。
• 整個流程用python用簡單的代碼實現，學習成本較低。
• 與pytorch和torchserve無縫銜接。

圖片

如上圖所示為torch-trt的優(yōu)化流程：

• Partition Graph（劃分圖） 首先要對 PyTorch 模型的計算圖進行分析，找出其中 TensorRT 所支持的節(jié)點。這是因為并非所有的 PyTorch 操作都能直接被 TensorRT 處理，需要確定哪些部分可以利用 TensorRT 的優(yōu)化能力。根據識別出的支持節(jié)點情況，決定哪些部分在 PyTorch 中運行，哪些部分可以在 TensorRT 中運行。對于可以在 TensorRT 中運行的部分，將進行后續(xù)的轉換操作。
• Compile TensorRT（編譯 TensorRT） 對于在 Partition Graph 步驟中確定可以由 TensorRT 處理的節(jié)點，將其轉換為 TensorRT 格式。這一步驟會利用 TensorRT 的優(yōu)化技術，如層融合、內核自動調整和內存優(yōu)化等，將這些節(jié)點轉換為高效的 TensorRT Engine（引擎），從而提高模型的推理速度。

加入torch-trt之后的優(yōu)化效果如上面表格，可以看出：

• trt優(yōu)化后，吞吐獲得了提升，符合預期，吞吐的提升主要是半精度和執(zhí)行流程優(yōu)化帶來的。顯存占用也因為半精度和算子融合大幅下降，符合預期。
• trt組的GPU沒有打滿，但是CPU打滿了，吞吐的瓶頸從GPU轉移到CPU，經過排查原因是預處理和后處理部分的CPU操作在請求量大的時候已經將CPU打滿了。這個問題過去是通過CPU微服務水平擴容解決的，在torchserve中CPU和GPU的執(zhí)行在一個進程內，無法水平擴容，下面將介紹解決方案。

3.2 預處理和后處理部分優(yōu)化

排查CPU執(zhí)行部分占用率較高的邏輯，原因為部分計算密集型邏輯被放在了前后處理中，例如opencv庫中的一些api執(zhí)行和通過numpy、pandas等庫進行矩陣計算。解決思路是將原來的python替換成NVIDIA官方提供的一些列對應的cuda版本庫。例如cvCuda和cuDF分別對應OpenCV與pandas，并且提供了相同的api，只需要在import包的時候進行替換，開發(fā)成本較低。

import cv2
import numpy as np
import cv2.cuda as cvcuda

# 讀取圖像
img = cv2.imread('your_image.jpg')

# 將圖像轉換為 GPU 上的格式
gpu_img = cvcuda.GpuMat(img)

# 使用 cvCuda 進行高斯模糊
gaussian_filter = cvcuda.createGaussianFilter(gpu_img.type(), -1, (5, 5), 1.5)
blurred_gpu = gaussian_filter.apply(gpu_img)

# 將處理后的圖像轉換回 CPU 格式
blurred_img = blurred_gpu.download()

# 顯示結果
cv2.imshow('Original Image', img)
cv2.imshow('Blurred Image (cvCuda)', blurred_img)
cv2.waitKey(0)
cv2.destroyAllWindows()

以上demo為例，只需要將cv2庫替換為cvCuda庫即可。此外計算密集型邏輯在CPU上執(zhí)行性價比并不高，很多時候是使用習慣所致，在深度學習時代之前OpenCV和pandas就比較流行了，后續(xù)沿用之前的用法。從下面NVIDIA官方提供的測試結果來看，相關計算邏輯在GPU上執(zhí)行可以獲得極大的性能提升。

圖片

上圖為OpenCV和CV-CUDA在不同算子上的吞吐表現對比。 

下圖展示了在同一個計算節(jié)點上（2x Intel Xeon Platinum 8168 CPUs , 1x NVIDIA A100 GPU），以 30fps 的幀率處理 1080p 視頻，采用不同的 CV 庫所能支持的最大的并行流數

圖片

下表為我們內部的測試結果：

從中可以看出,將預處理中計算密集型部分放在GPU上成功解決了此前遇到的問題。吞吐比base提升了4倍，GPU占用率提升顯著，CPU占用率沒有打滿。需要注意的是顯存占用率相應也會提升，因為之前在內存中進行的操作被移到了顯存上進行。

3.3 Torchserve on Kubernetes

torchserve官方通過Helm Charts提供了一個輕量級的k8s部署方案，可以實現服務的可靠運行和高可用性，，這也是我們在框架選型中看中torchserve的一個優(yōu)勢。

kubectl get pods

NAME                                             READY   STATUS    RESTARTS   AGE
grafana-cbd8775fd-6f8l5                          1/1     Running   0          4h12m
model-store-pod                                  1/1     Running   0          4h35m
prometheus-alertmanager-776df7bfb5-hpsp4         2/2     Running   0          4h42m
prometheus-kube-state-metrics-6df5d44568-zkcm2   1/1     Running   0          4h42m
prometheus-node-exporter-fvsd6                   1/1     Running   0          4h42m
prometheus-node-exporter-tmfh8                   1/1     Running   0          4h42m
prometheus-pushgateway-85948997f7-4s4bj          1/1     Running   0          4h42m
prometheus-server-f8677599b-xmjbt                2/2     Running   0          4h42m
torchserve-7d468f9894-fvmpj                      1/1     Running   0          4h33m

以上為一個集群的pods列表，除了torchserve的微服務階段，還提供了model-store功能，以及基于prometheus、grafana的監(jiān)控報警體系。再加上k8s原有的能力，既可實現一個輕量級的支持故障自動恢復、負載均衡、滾動更新、模型管理、安全配置的高可用性和可彈性擴展系統(tǒng)。

4  后續(xù)工作

torchserve在轉轉GPU推理服務系統(tǒng)里的落地是一次平衡開發(fā)效率與系統(tǒng)性能的工程實踐，總體上達到了計劃目標并且取得了業(yè)務價值。但是也存在一些不足，比如說原計劃前后處理部分線上和線下完全一致，共用相同的python代碼，但是實踐中遇到了CPU被打滿的情況改變了方案。盡管替代方案開發(fā)成本較低但是仍然做不到Write once, run anywhere。

此外，當前提供的解決方案以效率為先，兼顧性能，適合快速落地及迭代。后續(xù)計劃針對更復雜的業(yè)務場景(多模型推理等)及推理模型(llm推理等)，提供進階的解決方案。同時，在云原生平臺建設上，當前只是實現了一個入門版本，需要補課的內容還很多。

關于作者

楊訓政，轉轉算法工程方向架構師，負責搜索推薦、圖像、大模型等方向的算法工程架構工作。