pytorch模型推理之服务化

最新推荐文章于 2024-07-04 18:05:43 发布

RichardorMu

最新推荐文章于 2024-07-04 18:05:43 发布

阅读量1.6k

点赞数 1

分类专栏：深度学习文章标签： pytorch 模型部属

本文链接：https://blog.csdn.net/weixin_42264234/article/details/118637301

版权

深度学习专栏收录该内容

29 篇文章 7 订阅

订阅专栏

写在前面

因为项目的需要，开始接触模型推理（Model Inference）和模型服务化（Model Serving），即模型部署（Model Deployment）。近期在做PyTorch模型部署有些学习心得，趁热打铁记录下来。如果文章中有纰漏，非常欢迎斧正！

本文需要读者对torch有一定的了解，因为本文将讨论的重点是torch的模型保存的加载的办法、当前基于torch模型的几种服务化框架。

全文将介绍：

PyTorch模型的保存和加载，官方提供的主要有三种办法：保存模型参数（state_dict）、保存整个模型方法和保存成TorchScript。
PyTorch模型的服务化框架，笔者接触到的有两种：TorchServe[2]和Triton[3]。
PyTorch模型转化实战，主要介绍如何将模型转化为可方便服务化的TorchScript格式和ONNX通用格式。

1. PyTorch模型的保存和加载

只要谈到“深度学习”，“人工智能”，现在几乎所有人都会自动关联到TensorFlow和PyTorch，而当我们看到一个算法工程师的时候，往往会问：“你是TFboy还是PTboy（也有称TORCHboy）？”可见Tensorflow和PyTorch两种模型框架（Framework）在AI领域的地位。然而，当我们看一看TensorFlow和PyTorch的发展历史，却发现2017年才“问世”的PyTorch居然用了不到三年的时间就在市场占有率上和2015年的TensorFlow做到了“分庭抗礼”。笔者从博客《One simple graphic: Researchers love PyTorch and TensorFlow》[4]和文章《Pytorch vs Tensorflow in 2020》[5]中都给出了PyTorch和TensorFlow在2018、2019和2020的使用率
在这里插入图片描述

而且，在知乎话题《pytorch与tensorflow未来哪一个会占据更大的用户群体？》[6]中大部分答主都支持了PyTorch！

PyTorch既然这么受算法工程师的青睐，笔者作为一个纯工程师，为什么文章标题要说PyTorch难用？！就不怕被人喷么？

且把问题看清楚，PyTorch作为模型训练框架，在易用度、兼容性等方便肯定是比TensorFlow好的，这个结论是笔者结合线下调研和线上网友们的回复得到的。但是，诚如话题《pytorch与tensorflow未来哪一个会占据更大的用户群体？》高赞答主所述：在工业部署上，直接拿PyTorch的模型做服务化，实在太难用了！

那么，接下来，且看PyTorch模型保存和加载的几种方式，相信经过下面的介绍，不少读者就会明白难用在哪里了！

方法一：模型参数

这种方法是算法工程师最偏爱、也最受推崇的模型保存和加载的办法，具体模型保存和加载代码如下：

# 模型保存
# @model：PyTorch网络模型实例化对象
# @model_name：PyTorch模型名称
# @epoch：训练数据的迭代
torch.save(model.state_dict(), "{}_{}.pkl".format(model_name,epoch))

# 模型加载
# @model：PyTorch网络模型实例化对象
# @model_name：PyTorch模型名称
# @epoch：训练数据的迭代
model.load_state_dict(torch.load("{}_{}.pkl".format(model_name,epoch)))

当然，实际保存的模型名字的format是用户自己定义的，笔者只是给了自己常用的format。

从模型加载的API调用可以看到，这种方式实际保存的并不是模型，而仅仅是以字典形式（Dict）保存的模型参数而已；而且如果算法工程师把这个模型交给你，希望你帮忙做模型转化或者部署，那么你还得记得一定得把网络定义的class源码也得拿过来，否则你只能跟一堆kv键值对“大眼瞪小眼”。

就state_dict有更多兴趣、希望深入了解的读者欢迎阅读官网《SAVING AND LOADING MODELS》[7]。

方法二：整个模型

同样，我们来看这种方式的模型保存和加载的具体代码：

# 模型保存
# @model：PyTorch网络模型实例化对象
# @model_name：PyTorch模型名称
torch.save(model, "{}.pt".format(model_name))

# 模型加载
# @model_name：PyTorch模型名称
model = torch.load("{}.pt".format(model_name))

这种也是一部分算法工程师可能会用到的模型加载和保存的办法，不过没有方法一更好，具体原因将在第三章解释。

在笔者还对PyTorch一无所知的时候，曾经看到这个方法，以为自己得到了“救星”：哎，这种方法不是把整个PyTorch模型都保存下来了么？那岂不是不需要依赖算法工程师的python源代码就能部署模型了？然后就被现实打了脸。具体原因，同样的会在第三章给出解释。

方法三：TrochScript

先来看代码：

# 模型保存
# @model：PyTorch网络模型实例化对象
# @dummpy_input: PyTorch模型的输入Tensor
# @model_name：PyTorch模型名称
script_model = torch.jit.trace(model, dummpy_input)
torch.jit.sve(script_model, "{}.pth".format(model_name))

# 模型加载
# @model_name：PyTorch模型名称
script_model = torch.jit.load("{}.pth".format(model_name))

这种模型的保存和加载方法，很多算法工程师并没有接触过，甚至都不知道还有这种格式的存在！TorchScript是PyTorch 1.0.0版本才开始出现的，目的就是为了把模型静态化并保存：

TorchScript is a way to create serializable and optimizable models
from PyTorch code. Any TorchScript program can be saved from a Python
process and loaded in a process where there is no Python dependency.
TorchScript是一种从PyTorch代码创建可序列化和可优化模型的方法。任何TorchScript都可以从Python程序中保存，并在没有Python环境的程序中加载。
We provide tools to incrementally transition a model from a pure
Python program to a TorchScript program that can be run independently
from Python, such as in a standalone C++ program. This makes it
possible to train models in PyTorch using familiar tools in Python and
then export the model via TorchScript to a production environment
where Python programs may be disadvantageous for performance and
multi-threading reasons.
我们提供了将模型从纯Python程序逐步过渡到可以独立于Python运行的TorchScript程序的工具，例如在独立的C
++程序中。这样就可以使用Python中熟悉的工具在PyTorch中训练模型，然后通过TorchScript将模型导出到生产环境中，在该生产环境中Python程序可能由于性能和多线程原因而处于不利地位。

总结来说，PyTroch官方推出TorchScript格式就是为了做模型部署。

题话外：PyTorch模型的后缀名，无论是’.pt’、’.pth’还是’.pkl’甚至是’pb’，并不能绝对代表它具体是哪一类PyTorch模型。最有效的判断手段，还是得了解模型是怎么保存下来的；第三章还会介绍另一种办法。

2. PyTorch模型服务化框架

如果问到TensorFlow模型的部署/服务化框架，相信不少人听说过TensorFlow Serving[2]。但是如果问到PyTorch模型的部署/服务化框架呢？相信很多同学一时想不起来！

但其实，在2020年04月，AWS 和Facebook合作推出了PyTorch通用化模型部署/服务化框架TorchServe[2]。但是这个框架好用么？这里我个人非常赞同知乎话题《如何评价 PyTorch 在 2020 年 4 月推出的 TorchServe？》的高赞回答：显然不好用！而且PyTorch目前的三种格式做部署/服务化都不方便，甚至是为此而生的TorchScript格式！

而笔者上一系列介绍到Nvidia的Triton Inference Server[3]，官方文档中说也能支持PyTorch模型服务化，不过只能支持TorchScript格式。并且呢，需要用户提供一份带有input和output详细信息（包括名称、维度和数据类型）的配置文档。

总结TorchServe和Triton Inference Server的做法，大致可以分为：

对于非TorchScript模型来说，需要用户先转化为TorchScript模型；
对于TorchScript模型来说，则通过告知input和output相关信息，就可以做推理部署/服务化了。

笔者没有使用过TrorchServe，单单看了网上一些博主们介绍的文章，大多是通过TorchServe提供的demo来介绍怎么使用TorchServe，个人还是存疑：

首先，TorchScript提供了torch-model-archiver工具来帮助用户把非TorchScript模型转化到TorchScript模型。但是从用法来看，似乎只需要用户指定模型、模型网络定义的python文件即可（这里只讨论必选项）。这里有几个问题：

网络定义的python文件内容有要求么？如果是多层甚至多个网络定义怎么办？torch-model-archiver如何识别哪个才是对应的网络定义？（难道都试一遍么…）
不需要用户指定输入数据的维度么？因为无论是模型还是网络定义文件，都无法获取输入信息，包括输入名称、输入数据维度和数据类型？（非常疑惑）
其次，TorchServe部署必须把模型打包成.mar的形式么？如果是，那么已经保存成TorchScript格式的模型，如何打包？不用torch-model-archiver工具打包成.mar能用么？

III. PyTorch模型转化实战
一个类似于torch-model-archiver功能的工具，叫torch2onnx，顾名思义，主要目的还是为了把PyTorch模型转化为标准、通用、可移植的onnx模型；此外，这个工具也可以转化为TorchScript模型。工具已经放到了PYPI上，欢迎下载使用。源代码在GitHub：https://github.com/ooooona/model_tools 上.
在这里插入图片描述

当前版本是0.0.7，torch2onnx的用法：

$ torch2onnx -h
usage: torch2onnx [-h] [--cuda CUDA] [--src-model-path SRC_MODEL_PATH]
                  [--dst-model-dir DST_MODEL_DIR]
                  [--dst-model-name DST_MODEL_NAME] [-o ONNX]
                  [-t TORCH_SCRIPT] [-f NET_FILE] [-c CLS_NAME]
                  [--func-name FUNC_NAME] [-b BATCH] [--input-name INPUT_NAME]
                  [--input-shape INPUT_SHAPE [INPUT_SHAPE ...]]
                  [--output-name OUTPUT_NAME]
                  [--output-shape OUTPUT_SHAPE [OUTPUT_SHAPE ...]]



PyTorch Convertor Tool

optional arguments:
  -h, --help            show this help message and exit
  --cuda CUDA           Using CUDA or not (default: True)
  --src-model-path SRC_MODEL_PATH
                        Source Model's filename, Notice should be absolut path
                        (default: '')
  --dst-model-dir DST_MODEL_DIR
                        Destination Model's directory (default: './')
  --dst-model-name DST_MODEL_NAME
                        Destination Model's filename, Notice without postfix
                        (default: 'dst_model')
  -o ONNX, --onnx ONNX  Saving the current Model as ONNX (default: True)
  -t TORCH_SCRIPT, --torch-script TORCH_SCRIPT
                        Saving the current Model as TORCH-SCRIPT (default:
                        True)
  -f NET_FILE, --net-file NET_FILE
                        Torch Net Class Name (default: '')
  -c CLS_NAME, --cls-name CLS_NAME
                        Torch Net Class Name (default: '')
  --func-name FUNC_NAME
                        Function Name of getting Net Class Instance (default:
                        '')
  -b BATCH, --batch BATCH
                        Input Name for Model (default: 1)
  --input-name INPUT_NAME
                        Input Name for Model (default: 'input')
  --input-shape INPUT_SHAPE [INPUT_SHAPE ...]
                        Input Shape for Model, Notice without Batch-Size!
                        (default: (1, 28, 28))
  --output-name OUTPUT_NAME
                        Output Name for Model (default: 'output')
  --output-shape OUTPUT_SHAPE [OUTPUT_SHAPE ...]
                        Output Shape for Model, Notice without Batch-Size!
                        (default: (10))

这里解释下：模型路径必须指定，即--src-model-path，torch2onnx会根据用户提供的模型来识别具体是PyTorch哪种格式的模型：

如果是TorchScript，那么需要用户指定输入数据的维度，即–input-shape（请注意这个是不带batch维度的！）；至于输入的名称默认是input，如果用户需要修改，则通过指定–input-name修改即可；同样的，输入数据的batch默认是1，用户可通过–batch修改。
如果是非TorchScript，还要求用户指定网络模型定义的文件和该定义的Class的类名，分别为–net-file和–cls-name。当然，输入数据的维度，即–input-shape（请注意这个是不带batch维度的！），也是需要指定的；至于输入的名称、输入的batch修改，使用办法一样。
目前，torch2onnx工具只能支持单输入和单输出，后期会支持多输入。

接下来，简单介绍PyTorch的非TorchScript模型转化为TorchScript模型的办法，在此之前，三个知识点必须要先了解：

首先，但凡用torch.save()方法保存，无论是StateDict还是整个模型，想要调用torch.load()加载的话，必须有正确的python环境（torch.load()的版本只能≥torch.save()的版本，否则可能存在不兼容的情况），并且必须该模型的网络定义类必须被import到当前环境中。是的！哪怕是保存整个模型，还是必须把模型定义import进来，原因是：torch.save()本质上是pickle操作，仅仅是把实例内容序列化存储到磁盘上而已，反序列化回来的时候需要知道数据定义！这也解释了第一章提到的为什么会被打脸。
其次，用torch.save()保存的StateDict和整个模型，正确被torch.load()回来之后，数据类型是不一样的。StateDict的类型是collections.OrderedDict，嗯，所以说StateDict本质就是字典，保存的是kv键值对；而整个模型的话，类型则是网络定义的Class。
最后，TorchScript格式的模型，必须通过torch.jit.load()方法加载，如果用torch.load()加载则会抛出异常。
由此，一个自动识别用户的PyTorch模型的格式，并且转化为TorchScript的思路就呼之欲出了：

首先通过torch.load()尝试加载，如果抛出异常，则尝试以TorchScript方式torch.jit.load()加载。加载成功则直接告诉用户已经是TorchScript格式，不需要转化；加载失败，则认为不是PyTorch格式的模型，告警。如果没有抛出异常，则跳到第2步。
接着，检查用户是否提供了网络定义的python文件和类名，如果没有，提示用户非TorchScript格式的模型必须提供；反之，则把用户的python文件内的网络定义类import进来。成功后，跳到第3步。
然后，再判断torch.load()得到的对象的类型，如果是collections.OrderedDict，那么则是StateDict，用python的反射机制实例化网络定义类，调用.load_state_dict()加载，成功后跳到第4步；如果是用户提供的网络定义类，则已经得到了模型，直接跳到第4步。
非TorchScript模型转TorchScript模型，只需要调用torch.jit.trace()过一遍模型的实例化即可，最后调用torch.jit.save()就得到了TorchScript。
这里有一点需要注意：并不是所有非TorchScript格式都能转化到TorchScript格式的，这点跟网络结构有关。

写在最后

最后，笔者在探索PyTorch部署和服务化的过程中，请教过很多算法工程师。发现每一个PyTorch使用者非常拥护PyTorch，训练如何如何好用、方便。但是一涉及到模型部署上线的问题，一部分算法同学并不关心；另一部分最后只说了一句“转成onnx吧”。

由此可见，PyTorch虽然占领了学术界大半壁江山，但是真的要走到工业界，还是有一个很大的GAP要解决；而这一块，私以为TensorFlow做的真的不错。希望PyTorch官方能够把TorchScript做的再完善些，能做到完全独立的模型可移植（Portable）！下次再回来看PyTorch的部署的时候，我们这些工程人员也能跟算法人员一样，说一句“PyTorch真香”！

最后的最后，如果有模型工程、云原生、微服务相关的技术话题，非常欢迎交流和讨论~

文献和链接
[1]. 腾讯云智能钛：https://cloud.tencent.com/product/tione

[2]. TorchServe的GitHub地址：https://github.com/pytorch/serve

[3]. Triton Inference Server的GitHub地址：https://github.com/triton-inference-server/server

[4]. 《One simple graphic: Researchers love PyTorch and TensorFlow》：https://gradientflow.com/one-simple-graphic-researchers-love-pytorch-and-tensorflow/

[5]. 《Pytorch vs Tensorflow in 2020》：https://morioh.com/p/05e780149e4b

[6]. 《pytorch与tensorflow未来哪一个会占据更大的用户群体？》：https://www.zhihu.com/question/322651220

[7]. PyTorch官方手册《SAVING AND LOADING MODELS》：https://pytorch.org/tutorials/beginner/saving_loading_models.html

[8]. TensorFlow Serving的GitHub地址：https://link.zhihu.com/?target=https%3A//github.com/tensorflow/serving