如何部署 PyTorch 模型

熬夜写代码的平头哥

于 2024-03-26 18:54:48 发布

阅读量420

点赞数 5

文章标签：深度学习

本文链接：https://blog.csdn.net/lwx666sl/article/details/137054502

版权

使用 PyTorch 训练好了模型，如何部署到生产环境提供服务呢？

有 Web 开发经验的小伙伴可能首先想到的是用 HTTP 框架（例如 Flask）进行简单的封装，做一个简单的 HTTP 接口来对外提供服务。但既然是生产，那还免不了考虑性能、扩展性、可运维性等因素。另外，做好这个服务还要求开发人员具备良好的 Web 后端开发技术栈。好在 PyTorch 已经给我们提供了一套统一的标准来实现这些，也就是接下来要介绍的开源工具：TorchServe。

什么是 TorchServe？

TorchServe 是 PyTorch 官方开发的开源工具，源码地址 https://github.com/pytorch/serve

官方给出的描述是：

A flexible and easy to use tool for serving PyTorch models

从描述中就可知道 TorchServe 是用来部署 PyTorch 模型的，而它的特点是「可扩展性」和「易用性」。

为什么用它？

理由很多，这里列出笔者认为比较重要的几点。

一、PyTorch 官方背书，并且从源码的提交上看项目的开发进展非常活跃。

大厂出品无脑更风就对了，不然为什么这么多人用 TensorFlow（狗头）。

二、隐藏工程实现细节，对数据科学家友好。

使用 PyTorch 的工程师、科学家可能不知道如何实现工程相关的功能，例如 RPC、RESTful API，但他们一定懂得预处理（Preprocessing）和 PyTorch 模型（torch.nn.Module）。而知道这些就足够了，工程上的事儿，交给TorchServe 来处理，而且它做得还不戳。人生苦短，不要重复造轮子。

三、制定了「标准」。

由于 TorchServe 将系统工程和模型分开了，所以模型接入到 TorchServe 需要一套统一的标准，而 TorchServe 将这套「标准」制定得很简单。我们甚至可以直接将训练或测试的代码的一部分截取下来，稍加改动即可完成接入操作。

除了以上还有很多其他的使用 TorchServe 的理由，TorchServe 为我们提供了丰富的功能，例如日志、多进程，动态注册模型等。感兴趣可以参考 TorchServe 的文档

https://github.com/pytorch/serve/blob/master/docs/README.mdgithub.com/pytorch/serve/blob/master/docs/README.md

下面是 TorchServe 的完整架构图：

如何使用？

安装参考：https://github.com/pytorch/serve/blob/master/README.md#install-torchserve-and-torch-model-archiver

下面简单描述一下用 TorchServe 部署模型的几个步骤，你也可以参考笔者的 MNIST 完整 demo，地址为：

https://github.com/louis-she/torch-serve-mnistgithub.com/louis-she/torch-serve-mnist

第一步，打包模型

使用 torch-model-archiver 命令来打包模型（该命令在安装完 TorchServe 后会自动获得）。

你需要准备两到三个文件：

checkpoint.pth.tar

从命名就应该知道，这就是我们在训练过程中通过 torch.save 获得的模型权重文件，注意该文件内容只能包含模型的权重。参考文件：https://github.com/louis-she/torch-serve-mnist/blob/main/mnist/mnist.pth.tar

2. model.py

该文件应该包含单个模型的类，该模型类应该可以使用 load_state_dict 来成功加载 checkpoint.pth.tar 提供的权重。参考文件 https://github.com/louis-she/torch-serve-mnist/blob/main/mnist/model.py

3. handler.py （可选）

如果希望加入一些自定义的 preprocessing 和 postprocessing，可以加入该文件。比如，对于视觉类任务，我们至少需要对输入图片进行一些 resize，归一化操作，还会对模型输出做一些后处理。这些都可以在 handler.py 中来定义。TorchServe 已经为我们写好了一些通用的 Handler：https://github.com/pytorch/serve/tree/master/ts/torch_handler

你可以在打包的时候直接使用他们，或是写自己的 handler.py，自定义的 Handler 可以继承自 TorchServe 的任意 Handler，笔者的 MNIST demo 就继承了 ts.torch_handler.image_classifier. ImageClassifier 。

准备好以上几个文件后，我们就可以使用 torch-model-archiver 打包，下面是 MNIST demo 的打包命令：

torch-model-archiver --model-name mnist --serialized-file mnist.pth.tar --model-file model.py --handler handler --runtime python3 --version 1.0

这里的参数都比较容易理解，但注意 --model-name 参数我们可以取一个有意义的名称（但最好不要加中文或空格），该参数会影响到我们之后的调用服务的 URL 链接。

调用过后我们将会获得一个 mnist.mar 文件，这就是我们打包好的文件。

第二步，注册模型

创建一个目录，名称为 model-store，将第一步打包好的 .mar 复制到这个目录中，然后我们就可以启动 TorchServe 服务器程序了：

torchserve --start --model-store $model-store-path

$model-store-path 替换为 model-store 的绝对路径即可。服务进程启动后，model-store 中的 mar 文件**不会**自动注册到 TorchServe 中，我们需要通过 TorchServe 提供的 Management API 来注册 model-store 中的 mar 文件。

API 的调用都将使用 curl 命令，如果你对该命令不熟悉，也可以使用 http GUI 客户端。

首先注册 mar 文件，这里依然以 MNIST 为例子：

curl -X POST "http://localhost:8081/models?url=mnist.mar"

注意 url= 后面的值必须是在 model-store 目录中有的文件。

然后为 mnist 分配工作进程，你可以根据 CPU 核心数来分配

curl -v -X PUT "http://localhost:8081/models/mnist?min_worker=1&max_worker=1&synchronous=true"

现在我们的 mnist.mar 就可以对外提供服务了，准备一张手写数字图片 test.png，

>>> curl "http://127.0.0.1:8080/predictions/mnist" -T test.png
{
  "8": 0.9467765688896179,
  "3": 0.023649968206882477,
  "5": 0.019438084214925766,
  "9": 0.008277446031570435,
  "2": 0.001520493533462286
}

这里注意 url 的格式为 /predictions/$model-name，因为我们之前指定的 model-name 是 mnist，所以链接为 /predictions/mnist

看到这里你应该已经猜到，我们甚至可以同时注册多个模型，只要将打包好的 mar 文件放入 model-store 目录中，然后调用 Management API 后即可。