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

大规模服务 PyTorch Lightning 模型的完整指南。

作者 Caleb Kaiser，Cortex Labs

纵观机器学习领域，一个主要趋势是越来越多项目专注于将软件工程原则应用于机器学习。例如，Cortex 重现了部署无服务器函数的体验，但用于推理管道。同样，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 Trainer 是一个抽象了样板训练代码（如训练和验证步骤）的类，它有一个内置的 save_checkpoint() 函数，可以将模型保存为 .ckpt 文件。要将模型保存为检查点，只需将以下代码添加到你的训练脚本中：

现在，在我们开始部署这个检查点之前，需要注意的是，虽然我一直说“PyTorch Lightning 模型”，但 PyTorch Lightning 是 PyTorch 的一个封装——该项目的 README 中明确写道“PyTorch Lightning 只是有组织的 PyTorch”。因此，导出的模型是一个普通的 PyTorch 模型，可以相应地部署。

有了保存的检查点，我们可以在 Cortex 中非常轻松地部署模型。如果你不熟悉 Cortex，可以在这里快速了解，但使用 Cortex 进行部署的简单概述如下：

• 我们用 Python 为模型编写预测 API（应用程序编程接口）

• 我们在 YAML 中定义 API 的基础设施和行为

• 我们通过 CLI 中的命令部署 API

我们的预测 API 将使用 Cortex 的 Python Predictor 类来定义一个 init() 函数以初始化 API 并加载模型，以及一个 predict() 函数以在查询时提供预测：

非常简单。我们复用了一些训练代码，添加了一点推理逻辑，就这样。需要注意的是，如果你将模型上传到 S3（推荐），你需要添加一些访问它的逻辑。

接下来，我们在 YAML 中配置基础设施：

同样简单。我们给 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：

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

我们的 YAML 看起来也相当相似：

我在这里添加了一个监控标志，只是为了展示配置有多么简单，还有一些 ONNX 特定的字段，但除此之外 YAML 是一样的。

最后，我们使用与之前相同的 $ cortex deploy 命令进行部署，我们的 ONNX API 就上线了。

3. 使用 Torchscript 的 JIT 编译器进行序列化

对于最终的部署，我们将把 PyTorch Lightning 模型导出为 Torchscript，并使用 PyTorch 的 JIT 编译器来提供服务。要导出模型，只需在训练脚本中添加以下内容：

用于此的 Python API 与普通 PyTorch 示例几乎相同：

YAML 保持不变，CLI 命令当然也保持一致。如果需要，我们实际上可以通过简单地用新脚本替换旧的 predictor.py 脚本，并再次运行 $ cortex deploy 来更新之前的 PyTorch API 以使用新模型：

Cortex 在此自动执行滚动更新，即启动新的 API 并与旧 API 交换，从而防止模型更新期间出现任何停机。

就是这样。现在你拥有了一个完全可操作的预测 API，用于实时推理，从 Torchscript 模型提供预测。

那么，你应该使用哪种方法？

这里显而易见的问题是哪种方法性能最佳。事实是，这里没有直接的答案，因为这取决于你的模型。

对于像 BERT 和 GPT-2 这样的 Transformer 模型，ONNX 可以提供令人难以置信的优化（我们在 CPU 上测量到吞吐量提高了 40 倍）。对于其他模型，Torchscript 可能比原版 PyTorch 表现更好——尽管这也存在一些注意事项，因为并非所有模型都能干净地导出到 Torchscript。

幸运的是，由于使用任何选项部署都非常容易，你可以并行测试所有三种方法，看看哪种对你的特定 API 表现最佳。

简介：Caleb Kaiser（@KaiserFrose）是 Cortex Labs 创始团队成员，负责维护 Cortex。

原文。经许可转载。

转载声明：本文为转载整理，原文版权归原作者及原站点所有。

原文来源：KDnuggets

原创地址：https://www.kdnuggets.com/2020/11/deploy-pytorch-lightning-models-production.html#comment-184728