【51CTO.com快譯】縱觀機器學習領域，一大趨勢是側生于將軟件工程原理運用于機器學習的項目數(shù)量激增。比如說，Cortex再現(xiàn)了部署無服務器函數(shù)的體驗，不過借助推理管道。與之相仿，DVC實現(xiàn)了現(xiàn)代版本控制和CI/CD管道，但面向機器學習。

[[351235]]

PyTorch Lightning有類似的理念，僅運用于訓練。框架提供了面向PyTorch的Python包裝器，讓數(shù)據科學家和工程師可以編寫干凈、易于管理和高性能的訓練代碼。

我們之所以構建整個部署平臺，一方面是由于我們討厭編寫樣板代碼，因此我們是PyTorch Lightning的忠實擁護者。本著這種精神，我寫了這篇指南，介紹將PyTorch Lightning模型部署到生產環(huán)境。在此過程中，我們將介紹導出PyTorch Lightning模型、加入到推理管道中的幾種不同方法。

部署PyTorch Lightning模型用于推理的各種方法

三種方法可以導出PyTorch Lightning模型來部署：

• 將模型另存為PyTorch檢查點
• 將模型轉換成ONNX
• 將模型導出到Torchscript

我們可以使用Cortex滿足這三種方法。

1.直接包裝和部署PyTorch Lightning模塊

從最簡單的方法開始，不妨部署一個沒有任何轉換步驟的PyTorch Lightning模型。

PyTorch Lightning Trainer是抽象樣板訓練代碼(想想訓練和驗證步驟)的一個類，它有內置的save_checkpoint()函數(shù)，可將模型另存為.ckpt文件。要將模型另存為檢查點，只需將該代碼添加到訓練腳本中：

圖1

現(xiàn)在，開始部署該檢查點之前，要特別注意的是，雖然我一直說“PyTorch Lightning模型”，但PyTorch Lightning是PyTorch的包裝器——該項目的README實際上寫著“PyTorch Lightning只是有組織的PyTorch”。因此，導出的模型是普通的PyTorch模型，可相應部署。

有了保存的檢查點，我們可以在Cortex中很輕松地部署模型。如果您不熟悉Cortex，可以在這里(https://docs.cortex.dev/)快速熟悉一下，但是Cortex部署過程的簡單概述如下：

• 我們使用Python為模型編寫預測API
• 我們使用YAML定義API基礎架構和行為
• 我們從CLI使用命令來部署API

我們的預測API將使用Cortex的Python Predictor類來定義init()函數(shù)，以初始化我們的API并加載模型，并使用predict()函數(shù)在查詢時進行預測： 

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 

同樣很簡單。我們?yōu)锳PI取名，告訴Cortex我們的預測AI是哪個，并分配一些CPU資源。

接下來，我們部署它：

請注意：我們還可以部署到由Cortex啟動并管理的集群上：

圖3

針對所有部署，Cortex都會對我們的API進行容器化處理，并將其公開為Web服務。針對云部署，Cortex配置負載均衡、自動擴展、監(jiān)測、更新和其他許多基礎架構功能。

就是這樣!現(xiàn)在我們有一個實際的Web API可根據要求處理模型預測。

2.通過ONNX Runtime導出到ONNX并部署

我們現(xiàn)已部署了一個普通的PyTorch檢查點，不妨讓情況復雜一些。

PyTorch Lightning最近添加了一個方便的抽象，用于將模型導出到ONNX(以前您可以使用PyTorch內置的轉換函數(shù)，不過它們需要多一點的樣板代碼)。要將模型導出到ONNX，只需將該代碼添加到訓練腳本中：

圖4

請注意，輸入樣本應模仿實際模型輸入的形狀。

一旦您導出了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"} 
view rawpredictor.py hosted with ❤ by GitHub 

基本上一樣。唯一的區(qū)別是，不是直接初始化模型，我們通過onnx_client來訪問它，這是Cortex為部署我們的模型而啟動的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)測標志，只為了表明配置有多容易;有一些ONNX特有的字段，不過除此之外是同樣的YAML。

最后，我們使用與之前一樣的$ cortex deploy命令來部署，我們的ONNX API處于活躍狀態(tài)。

3. 使用Torchscript的JIT編譯器來初始化

至于最后的部署，我們將把PyTorch Lightning模型導出到Torchscript，并使用PyTorch的 JIT編譯器來部署。要導出模型，只需將這部分添加到訓練腳本中：

圖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與之前一樣，當然CLI命令是一致的。如果我們想要，可以實際上更新之前的PyTorch API以使用新模型，只需把舊的predictor.py腳本換成新腳本，并再次運行$ cortex deploy：

圖6

Cortex在此處自動執(zhí)行滾動更新，新的API創(chuàng)建，然后與舊的API交換，因而防止模型更新間歇的任何停機。

就是這樣?，F(xiàn)在，您已經有了用于實時推理的完全可操作的預測API，可從Torchscript模型進行預測。

那么，您應使用哪種方法?

顯而易見的問題是哪種方法效果最好。事實上，這里沒有簡單的答案，這取決于您的模型。

針對BERT和GPT-2之類的Transformer模型，ONNX可以提供出色的優(yōu)化(我們測得CPU上的吞吐量提高了40倍)。至于其他模型，Torchscript的性能可能勝過普通PyTorch，不過也有一些地方要注意，因為并非所有模型都能干凈地導出到Torchscript。

幸好，您可以并行測試所有這三種方法，看看哪種方法最適合您的特定API。

原文標題：How to deploy PyTorch Lightning models to production，作者：Caleb Kaiser

【51CTO譯稿，合作站點轉載請注明原文譯者和出處為51CTO.com】