PyTorch的生态和模型部署

PyTorch的生态和模型部署

1. PyTorch生态

前几章，我们学习了PyTorch的基本使用、能够定义和修改自己的模型、常用的训练技巧和PyTorch的可视化。

PyTorch的强大，跟PyTorch的活跃的开源社区离不开关系，开源社区围绕PyTorch所生产的一些列工具包和程序，这些优秀的工具包极大地方便了PyTorch在特定领域的使用。比如：

• TorchVision、TorchVideo：用于图片和视频处理
• torchtext：用于自然语言处理
• PyTorch Geometric：图卷积网络

接下来，会介绍PyTorch针对某个领域（图像、视频、文本等）的其中某一具代表性的工具包进行详细介绍，主要包括工具包的作者、所在机构、数据预处理工具、数据扩增、常用模型结构的预定义、预训练模型权重、常用损失函数、常用评测指标、封装好的训练&测试模块及可视化工具。

1.1 torchvision

1.1.1 torchvision简介

" The torchvision package consists of popular datasets, model architectures, and common image transformations for computer vision. "

其包括如下库：

• torchvision.datasets：包含在计算机视觉中的常见数据集
• torchvision.models：具体的预训练模型查看链接，可分为Classification、Semantic Segmentation、Object Detection、instance Segmentation and Keypoint Detection、Video Classification
• torchvision.transforms：针对图片进行一些数据预处理，如放大和缩小、水平或垂直翻转等，其相关操作参考链接，其实战可参考链接。
• torchvision.io：提供了视频、图片和文件的IO操作的功能，包括读取、写入、编解码处理操作。
  • 注：不同版本之间的torchvision.io有较大变化
  • 除了read_video()等方法，torchvision.io提供了一个细粒度的视频API torchvision.VideoReader()，需要安装ffmpeg然后源码重新编译torchvision
  • 使用Video相关API，最好按照PyAV库
• torchvision.ops：提供了计算机视觉的特定操作，包括但不限于NMS、RolAlign、RolPool，可以参考这里
• torchvision.utils：提供了一些可视化的方法，如将若干图片拼接、可视化检测和分割等效果，具体参考链接。

1.2 PyTorchVideo

随着传播媒介和视频平台的发展，视频逐渐成伪新一代的主流媒体，使得视频相关的深度学习模型越来越受关注，但是针对视频的深度学习模型有许多缺点：

• 计算资源耗费更多，没有高质量的model zoo
• 数据集处理麻烦，没有很好的视频处理工具
• 随着多模态的流行，亟需一个工具处理其他模态
• 模型的部署优化问题

PyTorchVideo因此而生，提供了加速视频理解研究所需的可重用、模块化和高效的组件，它支持不同的深度学习视频组件，如视频模型、视频数据集和视频特定转换。其组件如下图：

对应的官方链接

1.2.1 PyTorch的主要部件和亮点

PyTorchVideo提供了model zoo，使得可以使用各种先进的预训练模型及其评判基准，其亮点如下：

• 基于PyTorch：使用PyTorch构建，使得PyTorch生态系统组件的使用变得简单
• Model Zoo：提供了包含I3D、R(2+1)D、SlowFast、X3D、MViT等SOTA模型的高质量model zoo，并且其model zoo与PyTorch Hub做了整合
• 数据预处理和常见数据：支持Kinetics-400, Something-Something V2, Charades, Ava (v2.2), Epic Kitchen, HMDB51, UCF101, Domsev等主流数据集和相应的数据预处理，同时还支持randaug, augmix等数据增强trick。
• 模块化设计：提供许多模块方便用户调用修改，在PyTorchVideo中具体来说包括data, transforms, layer, model, accelerator等模块，方便用户进行调用和读取。
• 支持多模态：对多模态的支持包括了visual和audio，未来会支持更多模态，为多模态模型的发展提供支持。
• 移动端部署优化：支持针对移动端模型的部署优化（使用前述的PyTorchVideo/accelerator模块），模型经过PyTorchVideo优化了最高达7倍的提速，并实现了第一个能实时跑在手机端的X3D模型（实验中可以实时跑在2018年的三星Galaxy S8上，具体请见Android Demo APP）。

其安装可以使用pip，对应命令pip install pytorchvideo，需满足如下版本约束：

• python >= 3.7
• PyTorch >= 1.8.0, torchvision 同样需匹配
• CUDA >= 10.2
• ioPath：具体情况
• fvcore >= 0.1.4：具体情况

其model zoo和benchmark可参考链接

1.2.2 使用PyTorchVideo model zoo

提供了三种使用方法，并给了对应的tutorial

• TorchHub：tutorial
• PySlowFast：使用PySlowFast workflow训练或测试PyTorchVideo models和datasets
• PyTorch Lightning建立了一个工作流处理，参考官方tutorial
• 更多的教材，可参考链接

from torchinfo import summary
import torchvision.models as models

resnet = models.resnet18()
summary(resnet, (1, 3, 224, 224))

其中，可以看到torchinfo提供了更加详细的信息，包括每个模块的信息（每一层的信息，输出shape和参数量）、模型总共参数量、模型大小、前向/反向传播需要的内存，参数内存、预估总共信息。

1.3 torchtext

PyTorch官方用于自然语言处理NLP的工具包为torchtext，由于NLP和CV在数据处理中的布通，因此NLP的工具包torchtext和torchvision等CV相关工具功能上也有差异，如：

• 数据集定义方式布通
• 数据预处理工具
• 没有过多的model zoo

更多可参考链接，torchtext官方文档

1.3.1 torchtext的主要组成

torchtext可以方便地对文本进行预处理，如截断补长、构建词表等，其主要组成部分：

• 数据处理工具：torchtext.data.functional, torchtext.data.utils
• 数据集：torchtext.data.datasets
• 词表工具：torchtext.vocab
• 评测指标：torchtext.metrics

1.3.2 安装及使用

安装可以使用pip，命令为pip install torchtext

其使用主要如下：

• Field：torchtext中定义数据类型及转换为张量的指令，其明确了如何处理不同类型的数据，具体的处理由Dataset完成
• 词汇表（vocab）：在NLP中，将字符串形式的词语（word）转换为数字形式的向量表示（embedding）非常重要，被称为Word Embedding，其基本思想是收集一个比较大的语料库（与所做任务相关），使用word2vec等方法构建词语到向量的映射，可以使用Field自带的build_vocab函数完成，TEXT.build_vocab(train)
• 数据迭代器：即torchtext中的DataLoader
• 使用自带数据集：参考官方文档、

1.3.3 评测指标

NLP中部分任务的评测不是通过准确率等指标完成，比如机器翻译任务常用BLEU (bilingual evaluation understudy) score来评价预测文本和标签文本之间的相似程度。在torchtext可以调用torchtext.data.metrics.bleu_score来完成

1.3.4 Others

由于NLP常用的网络结构比较固定，其模型主要通过torch.nn模块实现，比如torch.nn.LSTM，torch.nn.RNN等

对于文本研究而言，当下Transformer已经成为了绝对的主流，因此PyTorch生态中的HuggingFace等工具包也受到了越来越广泛的关注。

2. PyTorch的模型部署

深度学习最终需要部署在手机端、开发板、嵌入式设备等终端，才能解决问题。但是通常情况下这些终端设备由于各种限制无法使用训练好的权重进行推理，因此需要将得到的权重变换从而部署到设备中，因此出现了以下模型部署pipeline：

这里将以ONNX为例介绍模型的部署，其流程通常如下：

• 使用torch.onnx进行模型格式转化
• 使用ONNX Runtime进行模型推理

2.1 ONNX和ONNX Runtime简介和安装

2.1.1 ONNX简介

• ONNX官网：https://onnx.ai/
• ONNX GitHub：https://github.com/onnx/onnx
• ONNX Tutorial：https://github.com/onnx/tutorials#converting-to-onnx-format

ONNX**( Open Neural Network Exchange)** 是 Facebook (现Meta) 和微软在2017年共同发布的，用于标准描述计算图的一种格式。各类设备只要基于兼容ONNX即可运行相关的模型。同时，各设备只要针对ONNX标准进行模型性能优化，就可以使得所有兼容ONNX的框架受益。目前ONNX支持如下深度学习框架和多种推理引擎：

2.1.2 ONNX Runtime简介：

• ONNX Runtime官网：https://www.onnxruntime.ai/
• ONNX Runtime GitHub：https://github.com/microsoft/onnxruntime

ONNX Runtime 是由微软维护的一个跨平台机器学习推理加速器，它直接对接ONNX，可以直接读取.onnx文件并实现推理，不需要再把 .onnx 格式的文件转换成其他格式的文件。

ONNX和ONNX Runtime的适配关系参考链接。

安装过程如下：

pip install onnx
pip install onnxruntime # use cpu
pip install onnxruntime-gpu # use gpu

ONNX Runtime和CUDA之间的适配关系参考链接。

2.2 PyTorch的模型部署

2.2.1 PyTorch模型导出为onnx

• 模型转换为ONNX格式

使用torch.onnx.export()把模型转换成 ONNX 格式的函数。模型导成onnx格式前，我们必须调用model.eval()或者model.train(False)以确保我们的模型处在推理模式下，避免因为dropout或batchnorm等运算符在推理和训练模式下的不同产生错误。

import torch.onnx 
# 转换的onnx格式的名称，文件后缀需为.onnx
onnx_file_name = "xxxxxx.onnx"
# 我们需要转换的模型，将torch_model设置为自己的模型
model = torch_model
# 加载权重，将model.pth转换为自己的模型权重
# 如果模型的权重是使用多卡训练出来，我们需要去除权重中多的module. 具体操作可以见5.4节
model = model.load_state_dict(torch.load("model.pth"))
# 导出模型前，必须调用model.eval()或者model.train(False)
model.eval()
# dummy_input就是一个输入的实例，仅提供输入shape、type等信息 
batch_size = 1 # 随机的取值，当设置dynamic_axes后影响不大
dummy_input = torch.randn(batch_size, 1, 224, 224, requires_grad=True) 
# 这组输入对应的模型输出
output = model(dummy_input)
# 导出模型
torch.onnx.export(model,        # 模型的名称
                  dummy_input,   # 一组实例化输入
                  onnx_file_name,   # 文件保存路径/名称
                  export_params=True,        #  如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False.
                  opset_version=10,          # ONNX 算子集的版本，当前已更新到15
                  do_constant_folding=True,  # 是否执行常量折叠优化
                  input_names = ['input'],   # 输入模型的张量的名称
                  output_names = ['output'], # 输出模型的张量的名称
                  # dynamic_axes将batch_size的维度指定为动态，
                  # 后续进行推理的数据可以与导出的dummy_input的batch_size不同
                  dynamic_axes={'input' : {0 : 'batch_size'},    
                                'output' : {0 : 'batch_size'}})

• 模型的验证（验证是否导出成功）

模型成功导出后得到一个ONNX格式的文件，可以使用onnx.checker.check_model()检测模型是否可用：

import onnx
# 我们可以使用异常处理的方法进行检验
try:
    # 当我们的模型不可用时，将会报出异常
    onnx.checker.check_model(self.onnx_model)
except onnx.checker.ValidationError as e:
    print("The model is invalid: %s"%e)
else:
    # 模型可用时，将不会报出异常，并会输出“The model is valid!”
    print("The model is valid!")

• ONNX可视化

ONNX格式的模型可以使用Netron可视化，下载链接。这里给出官网的squeezenet例子链接。

2.2.2 使用ONNX Runtime进行推理

完成模型的转化，检测后，即可使用ONNX Runtime运行转化后的模型查看推理结果。

# 导入onnxruntime
import onnxruntime
# 需要进行推理的onnx模型文件名称
onnx_file_name = "xxxxxx.onnx"

# onnxruntime.InferenceSession用于获取一个 ONNX Runtime 推理器
ort_session = onnxruntime.InferenceSession(onnx_file_name)  

# 构建字典的输入数据，字典的key需要与我们构建onnx模型时的input_names相同
# 输入的input_img 也需要改变为ndarray格式
ort_inputs = {'input': input_img} 
# 我们更建议使用下面这种方法,因为避免了手动输入key
# ort_inputs = {ort_session.get_inputs()[0].name:input_img}

# run是进行模型的推理，第一个参数为输出张量名的列表，一般情况可以设置为None
# 第二个参数为构建的输入值的字典
# 由于返回的结果被列表嵌套，因此我们需要进行[0]的索引
ort_output = ort_session.run(None,ort_inputs)[0]
# output = {ort_session.get_outputs()[0].name}
# ort_output = ort_session.run([output], ort_inputs)[0]

这里需注意：

• PyTorch模型的输入为Tensor，而ONNX输入为arry，因此需要实现张量到array的转化，通常如下：

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

• 输入的array的shape应该和导出模型的dummy_input的shape一致，不一致需要进行resize
• 得到的结果是一样list，需要进行索引才能得到array格式的结果
• 构建输入字典时，需要注意字典的key应与导出ONNX格式设置的input_name一致

2.2.3 Sample

这里给出官网的链接供参考。

References

1. 深入浅出PyTorch第八章PyTorch生态
2. 深入浅出PyTorch第九章模型部署