Torchserve在轉轉GPU推理架構中的實踐
1 背景
轉轉面向二手電商業務,在搜索推薦、智能質檢、智能客服等場景落地了AI技術。在實踐的過程中,也發現了存在GPU執行優化不充分,浪費計算資源,增加應用成本等問題。
此外還存在線上線下處理邏輯需要分別開發的情況,造成額外的開發成本和錯誤排查成本,對一些需要高速迭代的業務場景的負面影響不可忽視。本文將會重點介紹基于Torchserve進行推理服務部署架構優化的工程實踐,希望對面臨類似問題的同學們有所幫助。
2 問題和解決思路
2.1 現狀分析
圖片
上圖為之前的推理系統架構圖,采用CPU和GPU分離的架構。這種架構的特點是:
GPU部分和CPU部分解耦分別部署微服務,預處理部分一般是在CPU上執行,容易成為推理服務的性能瓶頸。解耦后可以將CPU部分部署在單獨的機器上,無限水平擴容擴容。
圖片
上圖為美團通用高效的推理服務部署架構方案,可以看到架構思路基本相同。
2.2 問題
該方案具有很大的優點,也是業界多個公司采取的方案,但是架構的選項也需要考慮具體的業務場景。該方案在轉轉的的場景里就出現了一些問題,比如:
- 迭代效率:GPU執行部分的是基于torch或者tf的模型,可以接近0成本的部署到推理微服務上,但是CPU執行部分包含一些預處理、后處理的部分,除了一些常見的圖像解碼、縮放、NMS等操作,還有很多自定義的邏輯,難以通過統一協議進行低成本部署,需要在微服務里進行二次開發,并且大部分情況下需要采用和離線算法實現不同的開發語言,拖累業務迭代效率。當前業界算法工程師的主要工作語言為python,對其他的語言或者dsl都存在一定的學習成本。
- 網絡通信:轉轉的部分業務場景存在圖片尺寸較大的情況,因為二手的某些質檢場景需要高清圖片判斷有沒有破損、劃痕、污漬等細小痕跡。這樣微服務之間的網絡通信開銷負擔比一般場景大很多。
2.3 解決思路
2.3.1框架調研
深度學習模型的常見的部署框架有以下幾個:
特性 | Triton | TorchServe | TensorFlow Serving |
支持的框架 | 多種深度學習框架,包括 TensorFlow、PyTorch、ONNX、TensorRT、OpenVINO 等 | 專為 PyTorch 設計,支持 PyTorch 模型 | 專為 TensorFlow 設計,支持 TensorFlow 模型 |
性能 | 高性能推理服務器,支持動態批處理、模型并行、多模型并發等 | 性能較好,支持多線程推理,GPU 支持良好 | 性能較好,支持多線程推理,GPU 支持良好 |
易用性 | 配置相對復雜,需要手動配置模型倉庫、推理服務等 | 易用性較好,提供了命令行工具和 Python API | 易用性較好,提供了命令行工具和 gRPC/REST API |
社區和支持 | 由 NVIDIA 開發和維護,社區活躍,文檔和示例豐富 | 由 Facebook 開發和維護,社區活躍,文檔和示例豐富 | 由 Google 開發和維護,社區非常活躍,文檔和示例豐富 |
從性能和質量的角度,三個框架水平都可以達到要求。重點考慮三個框架對于非深度學習部分的自定義邏輯支持:
- tensoflow的tf.function裝飾器和AutoGraph機制 tf.function是TensorFlow 2.x中引入的一個重要特性,它通過在Python函數上應用一個裝飾器,將原生Python代碼轉換為TensorFlow圖代碼,從而享受圖執行帶來的性能優勢。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等接口來完成模型的加載、推理和卸載。這種方式適合于需要復雜邏輯處理的場景,比如多模型協同工作或者需要自定義預處理和后處理的情況。Python Backend可以與Triton的pipeline功能結合,實現更復雜的推理流程。例如下面的代碼中,實現了一個initialize方法初始化,并且實現一個execute方法執行具體的邏輯,代碼為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在功能和機制上比較類似。都提供了一個靈活、易用且可擴展的解決方案,特別適合于需要快速部署和靈活處理模型的場景。框架兼容上triton支持主流的所有框架,torchserve主要支持pytorch和onnx協議,都可以滿足轉轉的需求。經過調研和試用,我們最終選擇了torchserve作為本次的框架選項,原因如下:
- torchserve與PyTorch生態深度集成,而Triton的學習曲線相對陡峭。torchserve主要支持torch框架,同時只兼容onnx協議。在輕量級和易用性上更符合轉轉當前的業務場景要求。例如在模型部署的格式轉換和配置上,torchserve相較于triton要簡易很多。
- 從后續轉轉GPU推理服務的演進來看,長期來看支持所有主流推理框架是必需的,短期在業務高速成長期優先選擇一個框架與長期支持所有主流并不沖突。
3 torchserve實踐過程
3.1 torchserve使用與調優
3.1.1 使用流程
圖片
以一個圖像模型簡單舉例,如圖所示:
- 將模型權重文件及前后處理邏輯python代碼打包成一個mar包
- mar包提交到torchserve進程中進行模型注冊
- 請求到來后執行圖片下載和模型前處理、推理、后處理,返回結果
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機制和易用性對于開發效率的提升是顯著的,在我們的內部場景里,一個單人維護的推理服務,在半年內節省了約32PD(人日)的開發成本。
3.1.2 torch-trt
模型的主干網絡部分,需要進行優化,否則執行效率差耗時較長。torch-trt允許將PyTorch 模型轉換為 TensorRT 格式,從而能夠利用 TensorRT 強大的優化引擎。TensorRT 針對 NVIDIA GPU 進行了高度優化,能夠實現快速的推理性能。它通過對模型進行層融合、內核自動調整和內存優化等操作,顯著提高了模型的推理速度。
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的優化流程:
- Partition Graph(劃分圖) 首先要對 PyTorch 模型的計算圖進行分析,找出其中 TensorRT 所支持的節點。這是因為并非所有的 PyTorch 操作都能直接被 TensorRT 處理,需要確定哪些部分可以利用 TensorRT 的優化能力。根據識別出的支持節點情況,決定哪些部分在 PyTorch 中運行,哪些部分可以在 TensorRT 中運行。對于可以在 TensorRT 中運行的部分,將進行后續的轉換操作。
- Compile TensorRT(編譯 TensorRT) 對于在 Partition Graph 步驟中確定可以由 TensorRT 處理的節點,將其轉換為 TensorRT 格式。這一步驟會利用 TensorRT 的優化技術,如層融合、內核自動調整和內存優化等,將這些節點轉換為高效的 TensorRT Engine(引擎),從而提高模型的推理速度。
分組 | GPU利用率 | CPU利用率 | QPS | 顯存占用 |
torch-base | 40% ~ 80% | 20%~40% | 10 | 2GB |
torch-trt | 10% ~ 50% | 100% | 17 | 680MB |
加入torch-trt之后的優化效果如上面表格,可以看出:
- trt優化后,吞吐獲得了提升,符合預期,吞吐的提升主要是半精度和執行流程優化帶來的。顯存占用也因為半精度和算子融合大幅下降,符合預期。
- trt組的GPU沒有打滿,但是CPU打滿了,吞吐的瓶頸從GPU轉移到CPU,經過排查原因是預處理和后處理部分的CPU操作在請求量大的時候已經將CPU打滿了。這個問題過去是通過CPU微服務水平擴容解決的,在torchserve中CPU和GPU的執行在一個進程內,無法水平擴容,下面將介紹解決方案。
3.2 預處理和后處理部分優化
排查CPU執行部分占用率較高的邏輯,原因為部分計算密集型邏輯被放在了前后處理中,例如opencv庫中的一些api執行和通過numpy、pandas等庫進行矩陣計算。解決思路是將原來的python替換成NVIDIA官方提供的一些列對應的cuda版本庫。例如cvCuda和cuDF分別對應OpenCV與pandas,并且提供了相同的api,只需要在import包的時候進行替換,開發成本較低。
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上執行性價比并不高,很多時候是使用習慣所致,在深度學習時代之前OpenCV和pandas就比較流行了,后續沿用之前的用法。從下面NVIDIA官方提供的測試結果來看,相關計算邏輯在GPU上執行可以獲得極大的性能提升。
圖片
上圖為OpenCV和CV-CUDA在不同算子上的吞吐表現對比。
下圖展示了在同一個計算節點上(2x Intel Xeon Platinum 8168 CPUs , 1x NVIDIA A100 GPU),以 30fps 的幀率處理 1080p 視頻,采用不同的 CV 庫所能支持的最大的并行流數
圖片
下表為我們內部的測試結果:
分組 | GPU利用率 | CPU利用率 | QPS | 顯存占用 |
torch-base | 40% ~ 80% | 20%~40% | 10 | 2GB |
torch-trt | 10% ~ 50% | 100% | 17 | 680MB |
GPU預處理+trt | 60% ~ 80% | 60% | 40 | 2.1GB |
從中可以看出,將預處理中計算密集型部分放在GPU上成功解決了此前遇到的問題。吞吐比base提升了4倍,GPU占用率提升顯著,CPU占用率沒有打滿。需要注意的是顯存占用率相應也會提升,因為之前在內存中進行的操作被移到了顯存上進行。
3.3 Torchserve on Kubernetes
torchserve官方通過Helm Charts提供了一個輕量級的k8s部署方案,可以實現服務的可靠運行和高可用性,,這也是我們在框架選型中看中torchserve的一個優勢。
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的監控報警體系。再加上k8s原有的能力,既可實現一個輕量級的支持故障自動恢復、負載均衡、滾動更新、模型管理、安全配置的高可用性和可彈性擴展系統。
4 后續工作
torchserve在轉轉GPU推理服務系統里的落地是一次平衡開發效率與系統性能的工程實踐,總體上達到了計劃目標并且取得了業務價值。但是也存在一些不足,比如說原計劃前后處理部分線上和線下完全一致,共用相同的python代碼,但是實踐中遇到了CPU被打滿的情況改變了方案。盡管替代方案開發成本較低但是仍然做不到Write once, run anywhere。
此外,當前提供的解決方案以效率為先,兼顧性能,適合快速落地及迭代。后續計劃針對更復雜的業務場景(多模型推理等)及推理模型(llm推理等),提供進階的解決方案。同時,在云原生平臺建設上,當前只是實現了一個入門版本,需要補課的內容還很多。
關于作者
楊訓政,轉轉算法工程方向架構師,負責搜索推薦、圖像、大模型等方向的算法工程架構工作。