如何將PyTorch Lighting模型部署到生產服務中

語言: CN / TW / HK

我報名參加金石計劃1期挑戰——瓜分10萬獎池,這是我的第15篇文章,點選檢視活動詳情

導讀

一篇用PyTorch Lighting提供模型服務的完全指南。

圖片

縱觀機器學習領域,一個主要趨勢是專注於將軟體工程原理應用於機器學習的專案。例如,Cortex重新創造了部署serverless功能的體驗,但使用了推理管道。類似地,DVC實現了現代版本控制和CI/CD管道,但是是針對ML的。\

PyTorch Lightning也有類似的理念,只適用於訓練。這些框架為PyTorch提供了一個Python包裝器,讓資料科學家和工程師可以編寫乾淨、易於管理和效能訓練的程式碼。

作為構建整個部署平臺的人,部分原因是我們討厭編寫樣板檔案,我們是PyTorch Lightning的忠實粉絲。本著這種精神,我整理了這篇將PyTorch Lightning模型部署到生產中的指南。在此過程中,我們將瞭解一些用於匯出PyTorch Lightning模型並將其包含在推理管道中的不同選項。

使用PyTorch Lightning模型進行推理的各種方法

有三種方法匯出用於PyTorch Lightning模型進行服務:

  • 儲存模型為PyTorch檢查點
  • 將模型轉換為ONNX
  • 匯出模型到Torchscript

我們可以用Cortex來對這三種進行服務。

1. 直接打包部署PyTorch Lightning模型

從最簡單的方法開始,讓我們部署一個不需要任何轉換步驟的PyTorch Lightning模型。

PyTorch Lightning訓練器是一個抽象了樣板訓練程式碼(想想訓練和驗證步驟)的類,它有一個內建的save_checkpoint()函式,可以將模型儲存為.ckpt檔案。要將你的模型儲存為一個檢查點,只需將以下程式碼新增到你的訓練指令碼中:

圖片

現在,在我們開始服務這個檢查點之前,重要的是要注意,當我一直說“PyTorch Lightning模型”時,PyTorch Lightning是PyTorch的一個封裝 —— 專案的自述檔案字面上說“PyTorch Lightning只是有組織的PyTorch”。因此,匯出的模型是一個正常的PyTorch模型,可以相應地提供服務。

有了儲存好的檢查點,我們就可以輕鬆地在Cortex中使用該模型。關於Cortex的部署過程的簡單概述如下:

  • 我們用Python為我們的模型編寫了一個預測API
  • 我們在YAML中定義api的基礎結構和行為
  • 我們通過CLI命令來部署API

我們的預測API將使用Cortex的Python Predictor類來定義一個init()函式來初始化我們的API並載入模型,以及一個predict()函式來在查詢時提供預測:

``` import torch import pytorch_lightning as pl import MyModel from training_code from transformers import (     AutoModelForSequenceClassification,     AutoConfig,     AutoTokenizer )

class PythonPredictor:     def init(self, config):         self.device = "cpu"         self.tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")         self.model = MyModel.load_from_checkpoint(checkpoint_path="./model.ckpt")

def predict(self, payload):         inputs = self.tokenizer.encode_plus(payload["text"], return_tensors="pt")         predictions = self.model(**inputs)[0]         if (predictions[0] > predictions[1]):           return {"class": "unacceptable"}         else:           return {"class": "acceptable"} ```

很簡單。我們從訓練程式碼中重新定義了一些程式碼,添加了一些推理邏輯,就是這樣。需要注意的一點是,如果你將模型上傳到S3(推薦),你需要新增一些訪問它的邏輯。

接下來,我們在YAML中配置基礎的設定:

- name: acceptability-analyzer   kind: RealtimeAPI   predictor:     type: python     path: predictor.py   compute:     cpu: 1

同樣也很簡單。我們給API一個名稱,告訴Cortex我們的預測API在哪裡,並分配CPU。

接下來,我們部署它:

圖片

注意,我們也可以將其部署到一個叢集中,並由Cortex進行管理:

圖片

在所有的部署中,Cortex將我們的API打包並將其作為web的服務公開。通過雲部署,Cortex可以配置負載平衡、自動擴充套件、監控、更新和許多其他基礎設施功能。

現在,我們有了一個實時的web API,可以通過請求用模型進行預測。

2. 匯出為ONNX並通過ONNX Runtime進行服務

現在我們已經部署了一個普通的PyTorch檢查點,讓我們把事情複雜化一點。

PyTorch Lightning最近添加了一個方便的抽象,用於將模型匯出到ONNX(以前,你可以使用PyTorch的內建轉換函式,儘管它們需要更多的樣板檔案)。要將模型匯出到ONNX,只需將以下程式碼新增到訓練指令碼中:

圖片

注意,輸入樣本應該模擬實際模型輸入的形狀。

一旦你匯出了ONNX模型,你就可以使用Cortex的ONNX Predictor來服務它。程式碼基本上是一樣的,過程也是一樣的。例如,這是一個ONNX預測API:

``` import pytorch_lightning as pl from transformers import (     AutoModelForSequenceClassification,     AutoConfig,     AutoTokenizer )

class ONNXPredictor:     def init(self, onnx_client, config):         self.device = "cpu"         self.client = onnx_client         self.tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")              def predict(self, payload):         inputs = self.tokenizer.encode_plus(payload["text"], return_tensors="pt")         predictions = self.client.predict(**inputs)[0]         if (predictions[0] > predictions[1]):           return {"class": "unacceptable"}         else:           return {"class": "acceptable"} ```

基本上是一樣的。唯一的區別是,我們不是直接初始化模型,而是通過onnx_client訪問它,這是一個ONNX執行時容器,Cortex為我們的模型提供服務。

我們的YAML看起來也很相似:

- name: acceptability-analyzer   kind: RealtimeAPI   predictor:     type: onnx     path: predictor.py     model_path: s3://your-bucket/model.onnx   monitoring:     model_type: classification

‍我在這裡添加了一個監視標誌,只是為了說明配置它是多麼容易,並且有一些ONNX特定的欄位,但除此之外是相同的YAML。

最後,我們使用與之前相同的$ cortex deploy命令進行部署,我們的ONNX API也是可用的。

3. 使用 Torchscript’s JIT compiler序列化

對於最後的部署,我們把PyTorch Lightning模型匯出到Torchscript,並使用PyTorch的JIT編譯器提供服務。要匯出模型,只需將此新增到你的訓練指令碼:

圖片

這個的Python API與普通的PyTorch示例幾乎相同:

``` import torch from torch import jit from transformers import (     AutoModelForSequenceClassification,     AutoConfig,     AutoTokenizer )

class PythonPredictor:     def init(self, config):         self.device = "cpu"         self.tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")         self.model = jit.load("model.ts")

def predict(self, payload):         inputs = self.tokenizer.encode_plus(payload["text"], return_tensors="pt")         predictions = self.model(**inputs)[0]         if (predictions[0] > predictions[1]):           return {"class": "unacceptable"}         else:           return {"class": "acceptable"} ```

YAML與以前一樣,CLI命令當然是一致的。如果我們願意,我們可以通過簡單地用新的指令碼替換舊的predictor.py指令碼來更新之前的PyTorch API,並再次執行$ cortex deploy:

圖片

在這裡,Cortex會自動執行滾動更新,即啟動一個新的API,然後與舊API進行交換,從而防止模型更新之間的停機時間。

就是這樣。現在,你有了一個用於實時推斷的完全可操作的預測API,從Torchscript模型提供預測。

那麼,你會用哪種方法呢?

這裡明顯的問題是哪種方法效能最好。事實上,這裡沒有一個簡單的答案,因為它取決於你的模型。

對於BERT和GPT-2這樣的Transformer模型,ONNX可以提供令人難以置信的優化(我們測量了cpu吞吐量有40倍提高)。對於其他模型,Torchscript可能比vanilla PyTorch表現得更好 —— 儘管這也帶來了一些警告,因為並不是所有的模型都清晰地匯出到Torchscript。

幸運的是,使用任何選項都可以很容易地進行部署,因此可以並行測試這三種方法,看看哪種最適合你的特定API.‍

英文原文:https://towardsdatascience.com/how-to-deploy-pytorch-lightning-models-to-production-7e887d69109f