【51CTO.com快译】纵观机器学习领域，一大趋势是侧生于将软件工程原理运用于机器学习的项目数量激增。比如说，Cortex再现了部署无服务器函数的体验，不过借助推理管道。与之相仿，DVC实现了现代版本控制和CI/CD管道，但面向机器学习。
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 Trainer是抽象样板训练代码(想想训练和验证步骤)的一个类，它有内置的save_checkpoint()函数，可将模型另存为.ckpt文件。要将模型另存为检查点，只需将该代码添加到训练脚本中：
图1
现在，开始部署该检查点之前，要特别注意的是，虽然我一直说“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()函数，以初始化我们的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"}
- 1.
- 2.
- 3.
- 4.
- 5.
- 6.
- 7.
- 8.
- 9.
- 10.
- 11.
- 12.
- 13.
- 14.
- 15.
- 16.
- 17.
- 18.
- 19.
- 20.
很简单。我们使用训练代码改变了一些代码的用途，并增添了一点推理逻辑。要注意的一点是，如果您将模型上传到了S3(推荐)，要添加访问模型的一些逻辑。
下一步，我们使用YAML配置基础架构：
- name: acceptability-analyzer
kind: RealtimeAPI
predictor:
type: python
path: predictor.py
compute:
cpu: 1
- 1.
- 2.
- 3.
- 4.
- 5.
- 6.
- 7.
同样很简单。我们为API取名，告诉Cortex我们的预测AI是哪个，并分配一些CPU资源。
接下来，我们部署它：
请注意：我们还可以部署到由Cortex启动并管理的集群上：
图3
针对所有部署，Cortex都会对我们的API进行容器化处理，并将其公开为Web服务。针对云部署，Cortex配置负载均衡、自动扩展、监测、更新和其他许多基础架构功能。
就是这样!现在我们有一个实际的Web API可根据要求处理模型预测。
2.通过ONNX Runtime导出到ONNX并部署
我们现已部署了一个普通的PyTorch检查点，不妨让情况复杂一些。
PyTorch Lightning最近添加了一个方便的抽象，用于将模型导出到ONNX(以前您可以使用PyTorch内置的转换函数，不过它们需要多一点的样板代码)。要将模型导出到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
- 1.
- 2.
- 3.
- 4.
- 5.
- 6.
- 7.
- 8.
- 9.
- 10.
- 11.
- 12.
- 13.
- 14.
- 15.
- 16.
- 17.
- 18.
- 19.
- 20.
基本上一样。唯一的区别是，不是直接初始化模型，我们通过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
- 1.
- 2.
- 3.
- 4.
- 5.
- 6.
- 7.
- 8.
- 9.
我在这里添加了监测标志，只为了表明配置有多容易;有一些ONNX特有的字段，不过除此之外是同样的YAML。
最后，我们使用与之前一样的$ cortex deploy命令来部署，我们的ONNX API处于活跃状态。
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
- 1.
- 2.
- 3.
- 4.
- 5.
- 6.
- 7.
- 8.
- 9.
- 10.
- 11.
- 12.
- 13.
- 14.
- 15.
- 16.
- 17.
- 18.
- 19.
- 20.
YAML与之前一样，当然CLI命令是一致的。如果我们想要，可以实际上更新之前的PyTorch API以使用新模型，只需把旧的predictor.py脚本换成新脚本，并再次运行$ cortex deploy：
图6
Cortex在此处自动执行滚动更新，新的API创建，然后与旧的API交换，因而防止模型更新间歇的任何停机。
就是这样。现在，您已经有了用于实时推理的完全可操作的预测API，可从Torchscript模型进行预测。
那么，您应使用哪种方法?
显而易见的问题是哪种方法效果最好。事实上，这里没有简单的答案，这取决于您的模型。
针对BERT和GPT-2之类的Transformer模型，ONNX可以提供出色的优化(我们测得CPU上的吞吐量提高了40倍)。至于其他模型，Torchscript的性能可能胜过普通PyTorch，不过也有一些地方要注意，因为并非所有模型都能干净地导出到Torchscript。
幸好，您可以并行测试所有这三种方法，看看哪种方法最适合您的特定API。
原文标题：How to deploy PyTorch Lightning models to production，作者：Caleb Kaiser
【51CTO译稿，合作站点转载请注明原文译者和出处为51CTO.com】
