如何將PyTorch Lightning模型部署到生產(chǎn)環(huán)境？
譯文【51CTO.com快譯】縱觀機(jī)器學(xué)習(xí)領(lǐng)域，一大趨勢(shì)是側(cè)生于將軟件工程原理運(yùn)用于機(jī)器學(xué)習(xí)的項(xiàng)目數(shù)量激增。比如說，Cortex再現(xiàn)了部署無服務(wù)器函數(shù)的體驗(yàn)，不過借助推理管道。與之相仿，DVC實(shí)現(xiàn)了現(xiàn)代版本控制和CI/CD管道，但面向機(jī)器學(xué)習(xí)。
PyTorch Lightning有類似的理念，僅運(yùn)用于訓(xùn)練。框架提供了面向PyTorch的Python包裝器，讓數(shù)據(jù)科學(xué)家和工程師可以編寫干凈、易于管理和高性能的訓(xùn)練代碼。
我們之所以構(gòu)建整個(gè)部署平臺(tái)，一方面是由于我們討厭編寫樣板代碼，因此我們是PyTorch Lightning的忠實(shí)擁護(hù)者。本著這種精神，我寫了這篇指南，介紹將PyTorch Lightning模型部署到生產(chǎn)環(huán)境。在此過程中，我們將介紹導(dǎo)出PyTorch Lightning模型、加入到推理管道中的幾種不同方法。
部署PyTorch Lightning模型用于推理的各種方法
三種方法可以導(dǎo)出PyTorch Lightning模型來部署：
- 將模型另存為PyTorch檢查點(diǎn)
- 將模型轉(zhuǎn)換成ONNX
- 將模型導(dǎo)出到Torchscript
我們可以使用Cortex滿足這三種方法。
1.直接包裝和部署PyTorch Lightning模塊
從最簡(jiǎn)單的方法開始，不妨部署一個(gè)沒有任何轉(zhuǎn)換步驟的PyTorch Lightning模型。
PyTorch Lightning Trainer是抽象樣板訓(xùn)練代碼(想想訓(xùn)練和驗(yàn)證步驟)的一個(gè)類，它有內(nèi)置的save_checkpoint()函數(shù)，可將模型另存為.ckpt文件。要將模型另存為檢查點(diǎn)，只需將該代碼添加到訓(xùn)練腳本中：
圖1
現(xiàn)在，開始部署該檢查點(diǎn)之前，要特別注意的是，雖然我一直說“PyTorch Lightning模型”，但PyTorch Lightning是PyTorch的包裝器——該項(xiàng)目的README實(shí)際上寫著“PyTorch Lightning只是有組織的PyTorch”。因此，導(dǎo)出的模型是普通的PyTorch模型，可相應(yīng)部署。
有了保存的檢查點(diǎn)，我們可以在Cortex中很輕松地部署模型。如果您不熟悉Cortex，可以在這里(https://docs.cortex.dev/)快速熟悉一下，但是Cortex部署過程的簡(jiǎn)單概述如下：
- 我們使用Python為模型編寫預(yù)測(cè)API
- 我們使用YAML定義API基礎(chǔ)架構(gòu)和行為
- 我們從CLI使用命令來部署API
我們的預(yù)測(cè)API將使用Cortex的Python Predictor類來定義init()函數(shù)，以初始化我們的API并加載模型，并使用predict()函數(shù)在查詢時(shí)進(jìn)行預(yù)測(cè)：
- 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"}
很簡(jiǎn)單。我們使用訓(xùn)練代碼改變了一些代碼的用途，并增添了一點(diǎn)推理邏輯。要注意的一點(diǎn)是，如果您將模型上傳到了S3(推薦)，要添加訪問模型的一些邏輯。
下一步，我們使用YAML配置基礎(chǔ)架構(gòu)：
- - name: acceptability-analyzer
- kind: RealtimeAPI
- predictor:
- type: python
- path: predictor.py
- compute:
- cpu: 1
同樣很簡(jiǎn)單。我們?yōu)锳PI取名，告訴Cortex我們的預(yù)測(cè)AI是哪個(gè)，并分配一些CPU資源。
接下來，我們部署它：
請(qǐng)注意：我們還可以部署到由Cortex啟動(dòng)并管理的集群上：
圖3
針對(duì)所有部署，Cortex都會(huì)對(duì)我們的API進(jìn)行容器化處理，并將其公開為Web服務(wù)。針對(duì)云部署，Cortex配置負(fù)載均衡、自動(dòng)擴(kuò)展、監(jiān)測(cè)、更新和其他許多基礎(chǔ)架構(gòu)功能。
就是這樣!現(xiàn)在我們有一個(gè)實(shí)際的Web API可根據(jù)要求處理模型預(yù)測(cè)。
2.通過ONNX Runtime導(dǎo)出到ONNX并部署
我們現(xiàn)已部署了一個(gè)普通的PyTorch檢查點(diǎn)，不妨讓情況復(fù)雜一些。
PyTorch Lightning最近添加了一個(gè)方便的抽象，用于將模型導(dǎo)出到ONNX(以前您可以使用PyTorch內(nèi)置的轉(zhuǎn)換函數(shù)，不過它們需要多一點(diǎn)的樣板代碼)。要將模型導(dǎo)出到ONNX，只需將該代碼添加到訓(xùn)練腳本中：
圖4
請(qǐng)注意，輸入樣本應(yīng)模仿實(shí)際模型輸入的形狀。
一旦您導(dǎo)出了ONNX模型，可以使用Cortex的ONNX Predictor來部署。代碼基本上看起來一樣，過程相同。比如說，這是一個(gè)ONNX預(yù)測(cè)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"}
- view rawpredictor.py hosted with ❤ by GitHub
基本上一樣。唯一的區(qū)別是，不是直接初始化模型，我們通過onnx_client來訪問它，這是Cortex為部署我們的模型而啟動(dòng)的ONNX Runtime容器。
我們的YAML看起來也很相似：
- - name: acceptability-analyzer
- kind: RealtimeAPI
- predictor:
- type: onnx
- path: predictor.py
- model_path: s3://your-bucket/model.onnx
- monitoring:
- model_type: classification
- view rawcortex.yaml hosted with ❤ by GitHub
我在這里添加了監(jiān)測(cè)標(biāo)志，只為了表明配置有多容易;有一些ONNX特有的字段，不過除此之外是同樣的YAML。
最后，我們使用與之前一樣的$ cortex deploy命令來部署，我們的ONNX API處于活躍狀態(tài)。
3. 使用Torchscript的JIT編譯器來初始化
至于最后的部署，我們將把PyTorch Lightning模型導(dǎo)出到Torchscript，并使用PyTorch的 JIT編譯器來部署。要導(dǎo)出模型，只需將這部分添加到訓(xùn)練腳本中：
圖5
這方面的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"}
- view rawpredictor.py hosted with ❤ by GitHub
YAML與之前一樣，當(dāng)然CLI命令是一致的。如果我們想要，可以實(shí)際上更新之前的PyTorch API以使用新模型，只需把舊的predictor.py腳本換成新腳本，并再次運(yùn)行$ cortex deploy：
圖6
Cortex在此處自動(dòng)執(zhí)行滾動(dòng)更新，新的API創(chuàng)建，然后與舊的API交換，因而防止模型更新間歇的任何停機(jī)。
就是這樣。現(xiàn)在，您已經(jīng)有了用于實(shí)時(shí)推理的完全可操作的預(yù)測(cè)API，可從Torchscript模型進(jìn)行預(yù)測(cè)。
那么，您應(yīng)使用哪種方法?
顯而易見的問題是哪種方法效果最好。事實(shí)上，這里沒有簡(jiǎn)單的答案，這取決于您的模型。
針對(duì)BERT和GPT-2之類的Transformer模型，ONNX可以提供出色的優(yōu)化(我們測(cè)得CPU上的吞吐量提高了40倍)。至于其他模型，Torchscript的性能可能勝過普通PyTorch，不過也有一些地方要注意，因?yàn)椴⒎撬心Ｐ投寄芨蓛舻貙?dǎo)出到Torchscript。
幸好，您可以并行測(cè)試所有這三種方法，看看哪種方法最適合您的特定API。
原文標(biāo)題：How to deploy PyTorch Lightning models to production，作者：Caleb Kaiser
【51CTO譯稿，合作站點(diǎn)轉(zhuǎn)載請(qǐng)注明原文譯者和出處為51CTO.com】
