如何将PyTorch Lighting模型部署到生产服务中


  磐创AI分享  

来源 | AI公园

作者 | Caleb Kaiser

编译 | ronghuaiyang

【导读】一篇用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

(0)

相关推荐