如何将PyTorch Lightning模型部署到生产环境?

编译作者: 布加迪 2020-11-10 08:00:00

【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()函数在查询时进行预测:

  1. import torch 
  2. import pytorch_lightning as pl 
  3. import MyModel from training_code 
  4. from transformers import ( 
  5.     AutoModelForSequenceClassification, 
  6.     AutoConfig, 
  7.     AutoTokenizer 
  8. class PythonPredictor: 
  9.     def __init__(self, config): 
  10.         self.device = "cpu" 
  11.         self.tokenizer = AutoTokenizer.from_pretrained("albert-base-v2"
  12.         self.model = MyModel.load_from_checkpoint(checkpoint_path="./model.ckpt"
  13.     def predict(self, payload): 
  14.         inputs = self.tokenizer.encode_plus(payload["text"], return_tensors="pt"
  15.         predictions = self.model(**inputs)[0] 
  16.         if (predictions[0] > predictions[1]): 
  17.           return {"class""unacceptable"
  18.         else
  19.           return {"class""acceptable"

很简单。我们使用训练代码改变了一些代码的用途,并增添了一点推理逻辑。要注意的一点是,如果您将模型上传到了S3(推荐),要添加访问模型的一些逻辑。

下一步,我们使用YAML配置基础架构:

  1. name: acceptability-analyzer 
  2.   kind: RealtimeAPI 
  3.   predictor: 
  4.     type: python 
  5.     path: predictor.py 
  6.   compute: 
  7.     cpu: 1 

同样很简单。我们为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:

  1. import pytorch_lightning as pl 
  2. from transformers import ( 
  3.     AutoModelForSequenceClassification, 
  4.     AutoConfig, 
  5.     AutoTokenizer 
  6. class ONNXPredictor: 
  7.     def __init__(self, onnx_client, config): 
  8.         self.device = "cpu" 
  9.         self.client = onnx_client 
  10.         self.tokenizer = AutoTokenizer.from_pretrained("albert-base-v2"
  11.          
  12.     def predict(self, payload): 
  13.         inputs = self.tokenizer.encode_plus(payload["text"], return_tensors="pt"
  14.         predictions = self.client.predict(**inputs)[0] 
  15.         if (predictions[0] > predictions[1]): 
  16.           return {"class""unacceptable"
  17.         else
  18.           return {"class""acceptable"
  19. view rawpredictor.py hosted with ❤ by GitHub 

基本上一样。唯一的区别是,不是直接初始化模型,我们通过onnx_client来访问它,这是Cortex为部署我们的模型而启动的ONNX Runtime容器。

我们的YAML看起来也很相似:

  1. name: acceptability-analyzer 
  2.   kind: RealtimeAPI 
  3.   predictor: 
  4.     type: onnx 
  5.     path: predictor.py 
  6.     model_path: s3://your-bucket/model.onnx 
  7.   monitoring: 
  8.     model_type: classification 
  9. view rawcortex.yaml hosted with ❤ by GitHub 

我在这里添加了监测标志,只为了表明配置有多容易;有一些ONNX特有的字段,不过除此之外是同样的YAML。

最后,我们使用与之前一样的$ cortex deploy命令来部署,我们的ONNX API处于活跃状态。

3. 使用Torchscript的JIT编译器来初始化

至于最后的部署,我们将把PyTorch Lightning模型导出到Torchscript,并使用PyTorch的 JIT编译器来部署。要导出模型,只需将这部分添加到训练脚本中:

图5

这方面的Python API与普通PyTorch示例一样:

  1. import torch 
  2. from torch import jit 
  3. from transformers import ( 
  4.     AutoModelForSequenceClassification, 
  5.     AutoConfig, 
  6.     AutoTokenizer 
  7. class PythonPredictor: 
  8.     def __init__(self, config): 
  9.         self.device = "cpu" 
  10.         self.tokenizer = AutoTokenizer.from_pretrained("albert-base-v2"
  11.         self.model = jit.load("model.ts"
  12.     def predict(self, payload): 
  13.         inputs = self.tokenizer.encode_plus(payload["text"], return_tensors="pt"
  14.         predictions = self.model(**inputs)[0] 
  15.         if (predictions[0] > predictions[1]): 
  16.           return {"class""unacceptable"
  17.         else
  18.           return {"class""acceptable"
  19. view rawpredictor.py hosted with ❤ by GitHub 

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】

机器学习 技术 科学
上一篇:联想持续发力智能物联,构建新基建时代下的行业新引擎 下一篇:图像也能做情感迁移?罗切斯特大学团队提出计算机视觉新任务
评论
取消
暂无评论,快去成为第一个评论的人吧

更多资讯推荐

什么是边缘人工智能和边缘计算?

边缘人工智能(Edge AI)是人工智能领域引人注目的新领域之一,其目的是让用户运行人工智能流程而不必担心隐私或数据传输较慢带来的影响。边缘人工智能可以使人工智能技术得到更广泛的应用,使智能设备在无需接入云平台的情况下对输入做出快速反应。

李睿 ·  1天前
第四范式NeurIPS 2020:知识图谱嵌入的自动化

前不久,来自第四范式的资深研究员姚权铭博士和大家分享了其参与并被 NeurIPS 2020 接收的论文《Interstellar: Searching Recurrent Architecture for Knowledge Graph Embedding》。

机器之心 ·  1天前
AI的偏见:机器就是绝对理性的么?

人工智能已经遍布我们的日常生活。从YouTube的首页推荐到生产药物,它无所不在,它对我们生活的影响超出我们的想象。但人工智能一定公平吗?不,绝对不是。

读芯术 ·  3天前
机器学习平台在Kubernetes上的实践

本文回顾总结近一段时间网易云音乐机器学习平台(GoblinLab)在容器化实践的一些尝试。截止今日,音乐机器学习平台(GoblinLab)在容器化方面的尝试,已开展了一段时间,并且已经有了阶段性的成果。

王军正 ·  4天前
盘点:2020年最酷的12家机器学习初创公司

我们筛选了12家初创公司,有的已经成立了几年时间,有的刚刚起步,但他们都在致力于解决与机器学习相关的各种挑战。

CRN ·  4天前
3D视觉CV界的终极体现形式,计算机如何「看」这个三维世界

现如今,随处可见 3D 视觉技术带来的便利,工业机器人、工件识别与定位、3D 成像技术、产品虚拟设计、智能制造、自动驾驶、SLAM、无人机、3D 重建、人脸识别等等,都涉及到 3D 视觉相关内容。

佚名 ·  4天前
人脸识别技术法律缺口亟待补上

由于进入门槛不高、成本低等原因,几乎任何单位、任何场所都可以用,这就造成了风险的不确定性。

社论 ·  2020-11-26 10:04:25
微分方程VS机器学习,实例讲解二者异同

微分方程与机器学习作为 AI 领域建模的两种方法,各自有什么优势?

佚名 ·  2020-11-25 14:40:48
Copyright©2005-2020 51CTO.COM 版权所有 未经许可 请勿转载