得物 AI 平台 -KubeAI 推理训练引擎设计和实践
1. KubeAI 介绍
KubeAI 是得物 AI 平台，是我们在容器化过程中，逐步收集和挖掘公司各业务域在 AI 模型研究和生产迭代过程中的需求，逐步建设而成的一个云原生 AI 平台。KubeAI 以模型为主线提供了从模型开发，到模型训练，再到推理(模型)服务管理，以及模型版本持续迭代的整个生命周期内的解决方案。
在数据方面，KubeAI 提供基于 cvat 的标注工具，与数据处理及模型训练流程打通，助力线上模型快速迭代；提供任务/Pipeline 编排功能，对接 ODPS/NAS/CPFS/OSS 数据源，为用户提供一站式 AI 工作站。平台自研推理引擎助力业务在提高模型服务性能的同时还能控制成本；自研训练引擎提高了模型训练任务吞吐量，缩短了模型的训练时长，帮助模型开发者加速模型迭代。
此外，随着 AIGC 的火热发展，我们经过调研公司内部 AI 辅助生产相关需求，上线了 AI 制图功能，为得物海报、营销活动、设计师团队等业务场景提供了基础能力和通用 AI 制图能力。
此前，我们通过一文读懂得物云原生AI平台-KubeAI的落地实践过程一文，向大家介绍了 KubeAI 的建设和在业务中的落地过程。本文，我们将重点介绍下 KubeAI 平台在推理、训练和模型迭代过程中的核心引擎能力实践经验。****
2.AI 推理引擎设计实现
2.1 推理服务现状及性能瓶颈分析
Python 语言以其灵活轻盈的特点，以及其在神经网络训练与推理领域提供了丰富的库支持，在模型研究和开发领域被广泛使用，所以模型推理服务也主要以 Python GPU 推理为主。模型推理过程一般涉及预处理、模型推理、后处理过程，单体进程的方式下 CPU 前/后处理过程，与 GPU 推理过程需要串行，或者假并行的方式进行工作，大致流程如下图所示：
上述架构的优势是代码写起来比较通俗易懂，但在性能上有很大的弊端，所能承载的 QPS 比较低。通过在 CV 域的模型上进行压测，我们发现推理 QPS 很难达到 5，深入分析发现造成这一问题的原因如下：
（1）单线程模式下，CPU 逻辑与 GPU 逻辑相互等待，GPU Kernel 函数调度不足，导致 GPU 使用率不高，无法充分提升服务 QPS。这种情况下只能开启更多进程来提升 QPS，但是更多进程会带来更大的 GPU 显存开销。
（2）多线程模式下，由于 Python 的 GIL 锁的原因，Python 的多线程实际上是伪的多线程，并不是真正的并发执行，而是多个线程通过争抢 GIL 锁来执行，这种情况下 GPU Kernel Launch 线程不能得到充分的调度。此外，在 Python 推理服务中开启多线程反而会导致 GPU Kernel Launch 线程频繁被 CPU 的线程打断，所以 GPU 算力也会一直“萎靡不振”，持续低下。
以上问题使得 如果推理服务想要支撑更多的流量，只能做横向的增加服务实例数，伴随着成本的上涨。
2.2 自研推理服务统一框架 kubeai-inference-framework
针对以上问题，KubeAI 的解决方案是把 CPU 逻辑与 GPU 逻辑分离在两个不同的进程中：CPU 进程主要负责图片的前处理与后处理，GPU 进程则主要负责执行 CUDA Kernel 函数，即模型推理。
为了方便模型开发者更快速地接入我们的优化方案，我们基于 Python 开发了一个 CPU 与 GPU 进程分离的统一框架***kubeai-inference-framework***，旧有 Flask 或 Kserve 的服务，稍作修改即可接入推理引擎统一框架，新增服务按照框架实现指定 function 即可。推理服务统一框架构如下图所示：
如前所述，推理服务统一框架的主要思路是把 GPU 逻辑与 CPU 逻辑分离到两个进程，除此之外，还会拉起一个 Proxy 进程做路由转发。
CPU 进程
CPU 进程主要负责推理服务中的 CPU 相关逻辑，包括前处理与后处理。前处理一般为图片解码，图片转换，后处理一般为推理结果判定等逻辑。CPU 进程在前处理结束后，会调用 GPU 进程进行推理，然后继续进行后处理相关逻辑。CPU 进程与 GPU 进程通过共享内存或网络进行通信，共享内存可以减少图片的网络传输。
GPU 进程
GPU 进程主要负责运行 GPU 推理相关的逻辑，它启动的时候会加载很多模型到显存，然后在收到 CPU 进程的推理请求后，直接触发 Kernel Lanuch 调用模型进行推理。
_kubeai-inference-framework_框架中对模型开发者提供了一个_Model_类接口，他们不需要关心后面的调用逻辑，只需要填充其中的前处理，后处理的业务逻辑，就可以快速上线模型服务，自动拉起这些进程。
Proxy 进程
Proxy 进程是推理服务入口，对外提供调用接口，负责路由分发与健康检查。当 Proxy 进程收到请求后，会轮询调用 CPU 进程，分发请求给 CPU 进程进行处理。
自研的推理服务统一框架，把 CPU 逻辑（图片解码，图片后处理等）与 GPU 逻辑（模型推理）分离到两个不同的进程中后，有效解决了 Python GIL 锁带来的 GPU Kernel Launch 调度问题，提升了 GPU 利用率，提高了推理服务性能。针对线上的某个推理服务，使用我们的框架进行了 CPU 与 GPU 进程分离，压测得出的数据如下表所示，可以看到 QPS 提升了近 7 倍。
2.3 做的更好 — 引入 TensorRT 优化加速
在支持推理服务接入_kubeai-inference-framework_统一框架的过程中，我们继续尝试在模型本身做优化提升。经过调研和验证，我们将现有 pth 格式模型通过转成 TensorRT 格式，并开启 FP16，在推理阶段取得了更好的 QPS 提升，最高可到 10 倍提升。
TensorRT 是由英伟达公司推出的一款用于高性能深度学习模型推理的软件开发工具包，可以把经过优化后的深度学习模型构建成推理服务部署在实际的生产环境中，并提供基于硬件级别的推理引擎性能优化。业内最常用的 TensorRT 优化流程，是把 pytorch / tensorflow 等模型先转成_onnx_格式，然后再将_onnx_格式转成 TensorRT（trt）格式进行优化，如下图所示：
TensorRT 所做的工作主要在两个时期，一个是网络构建期，另外一个是模型运行期。
网络构建期
模型解析与建立，加载 onnx 网络模型。
计算图优化，包括横向算子融合，或纵向算子融合等。
节点消除，去除无用的节点。
多精度支持，支持 FP32/FP16/int8 等精度。
基于特定硬件的相关优化。
模型运行期
序列化，加载 RensorRT 模型文件。
提供运行时的环境，包括对象生命周期管理，内存显存管理等
为了更好地帮助模型开发者使用 TensorRT 优化，KubeAI 平台提供了***kubeai-trt-helper*** 工具，用户可以使用该工具把模型转成 TensorRT 格式，如果在模型转换的过程中出现精度丢失等问题，也可以使用该工具进行问题定位与解决。_kubeai-trt-helper_主要在两个阶段为用户提供帮助：一个是问题定位，另一个阶段是模型转换。
问题定位
问题定位阶段主要是为了解决模型转 TensorRT 开启 FP16 模式时出现的精度丢失问题。一般分类模型，对精度的要求不是极致的情况下，尽量开启 FP16，FP16 模式下，NVIDIA 对于 FP16 有专门的 Tensor Cores 可以进行矩阵运算，相比 FP32 来说吞吐量提升一倍以上。比如在转 TensorRT 时，开启 FP16 出现了精度丢失问题，_kubeai-trt-helper_工具在问题定位阶段的大致工作流程如下：
第 1 步：设定模型转换精度要求后，标记所有算子为输出，然后对比所有算子的输出精度。
第 2 步：找到最早的不符合精度要求的算子，对该算子进行如下几种方式干预。
标记该算子为 FP32。
标记其父类算子为 FP32。
更改该算子的优化策略。
循环通过以上 2 个步骤，最终找到符合目标精度要求的模型参数。这些参数比如：需要额外开启 FP32 的那些算子等。相关参数会输出到配置文件中，如下：
模型转换模型转换阶段则直接使用上面问题定位阶段得到的参数，调用 TensorRT 相关接口与工具进行转换。此外，我们在模型转换阶段，针对 TensorRT 原有参数与 API 过于复杂的问题也做了一些封装，提供了更为简洁的接口，比如工具可以自动解析 onnx，判断模型的输入与输出 shape，不需要用户再提供相关 shape 信息等。
2.4 落地实践成果
在实际应用中，我们帮助算法域的模型开发同学，能够对一个推理基于自研推理服务统一框架进行实现的同时，也开启 TensorRT 优化，这样往往可以得到 QPS 两次优化的叠加效果。
2.4.1 分类模型，CPU 与 GPU 分离，TensorRT 优化，并开启 FP16，得到 10 倍 QPS 提升
线上某个基于 Resnet 的分类模型，对精度损失可以接受误差在 0.001（误差定义：median,atol,rtol）范围内。因此我们对该推理服务进行了 3 项性能优化：
使用_kubeai-inference-framework_统一框架，对 CPU 进程和 GPU 进程进行分离改造。
对模型转 ONNX 后，转 TensorRT。
开启 FP16 模式，并使用自研工具定位到中间出现精度损失的算子，把这些算子标记为 FP32。
经过以上优化，最终得到了 10 倍 QPS 的提升（与原来 Pytorch 直接推理比较），服务成本大幅削减。
2.4.2 检测模型，CPU 与 GPU 分离，TensorRT 模型优化，QPS 提升 4-5 倍左右。
线上某个基于 Yolo 的检查模型，由于对精度要求比较高，所以不能开启 FP16，我们直接在 FP32 的模式下进行了 TensorRT 优化，并使用_kubeai-inference-framework_统一框架对 GPU 进程与 CPU 进程分离，最终得到 QPS 4-5 倍的提升。
2.4.3 模型推理进程多实例化，充分利用 GPU 算力资源
在实际的场景中，往往 GPU 的算力是充足的，而 GPU 显存是不够的。经过 TensorRT 优化后，模型运行时需要的显存大小一般会降低到原来的 1/3 到 1/2。所以为了充分利用 GPU 算力，_kubeai-inference-framework_统一框架进一步优化，支持可以把 GPU 进程在一个容器内复制多份，这种架构即保证了 CPU 可以提供充足的请求给 GPU，也保证了 GPU 算力充分利用。
线上某个模型，经过 TensorRT 优化后，显存由原来的 2.4G 降低到只需要 1.2G。在保持推理服务配置 5G 显存不变的情况下，我们将 GPU 进程为复制 4 份，充分利用了 5G 显存，使得服务吞吐达到了原来的 4 倍。
3.AI 训练引擎优化实践
3.1 PyTorch 框架概况
PyTorch 是近年来较为火爆的深度学习框架，几乎占据了 CV（Computer Vision，计算机视觉）、NLP（Natural Language Processing，自然语言处理）领域各业务方向，算法同学基本都在使用 PyTorch 框架来进行模型训练。下图是基于 PyTorch 框架进行模型训练时的代码基本流程：
第 1 步：从 pytorch dataloader 中将本 step 训练过程中需要的数据拉出来。
第 2 步：将获取到的数据，例如：样本图片、样本标签的 tensor 等数据，复制到 GPU 显存里。
第 3 步：开始正式的模型训练：前向计算、计算损失、计算梯度、 更新参数。
整个训练过程的耗时，也主要分布在上面 3 个步骤。通常第 2 步不会是瓶颈，因为大部分训练样本图片都是被 resize 变小之后才从内存拷贝到到 GPU 显存上的。但由于模型的差异性、训练数据的差异性，经常是第 1、2 步会在训练过程中出现性能瓶颈，导致训练耗时长，GPU 利用率低下，影响模型迭代效率。
3.2 Dataloader 瓶颈分析及优化
3.2.1 PyTorch Dataset/Dataloader 分析
PyTorch 训练读取数据部分主要是通过 Dataset、Dataloader 的方式完成的，其中 Dataset 为用户自定义读取数据的类（继承自 torch.utils.data.Dataset），而 Dataloader 是 PyTorch 实现的在训练过程中对 Dataset 的调度器。
参数解释如下：
dataset(Dataset)：传入的自定义 Dataset（数据读取的具体步骤）。
batch_size(int, optional)：每个 batch 有多少个样本，每个 iter 可以从 dataloader 中取出多少数据。
shuffle(bool, optional)：在每个 epoch 开始的时候，对数据进行重新排序，可以使每个 epoch 读取数据的组合和顺序不同。
num_workers (int, optional)：这个参数决定 dataloader 启动几个后台进程来做数据拉取。0 意味着所有的数据都会被 load 进主进程，默认为 0。
collate_fn (callable, optional)：将一个 list 的 sample 组成一个 mini-batch 的函数，一般 CV 场景是 concat 函数。
pin_memory (bool, optional)：如果设置为 True，那么 data loader 将会在返回 batch 之前，将 tensors 拷贝到 CUDA 中的固定内存（CUDA pinned memory）中， 这个参数某些场景下有妙用。
drop_last (bool, optional)：该参数是对最后的未完成的 batch 来说的，比如 batch_size 设置为 64，而一个 epoch 只有 100 个样本，如果设置为 True，那么训练的时候后面的 36 个就被扔掉了，否则会继续正常执行，只是最后的 batch_size 会小一点。默认设置为 False。
上述参数中，比较重要的是num_workers，Dataloader 在构造的时候，会启动num_workers个 worker 进程，然后主进程会向 worker 进程分发读取任务，worker 进程读到数据之后，再把数据放到队列中供主进程取用。多进程模式使用的是torch.multiprocessing接口，可以实现 worker 进程与主进程之间共享内存，而且共享内存中可以存放 tensor，这样进程中如果返回 tensor，可以通过共享内存的方式直接将结果返回给主进程，减少多进程间的通讯开销。
当num_workers 为 0 的时候： get_data()流程与train_model()过程是串行，效率非常低下，如下图所示：
当num_workers 大于 0 开启多进程读取数据， 并且读取一个 batch 数据的时间小于一个 step 训练的时间时效率最高，GPU 算力被充分利用，如下图所示：
当num_workers 大于 0 开启多进程度数据， 但是读取一个 batch 数据的时间大于一个 step 训练的时间时，会出现 GPU 训练过程等待数据拉取，就会出现 GPU 算力空闲，训练耗时增加，如下图所示：
由此可见 Dateset 中的__getitem__函数非常重要，详细分析它的源码实现后我们发现，该函数的耗时主要包含 2 段时间：
load_image_time：从磁盘或者远程盘上读取数据的耗时。
transform_image_time：将图片或文本数据进行预处理的耗时。
3.2.2 解问题 — 设置合理的参数很重要
通过上一小节的分析，训练时相关参数的选择至关重要。总结如下：
batch_size：根据数据量，以及期望训练时长，用户合理自定义设置
训练环境（KubeAI Notebook/任务/流水线节点）的 CPU 配置：建议 CPU 配置为 GPU 卡数*(单 GPU 卡配置的 CPU 核数)。
num_workers：参数最小设置为 训练环境的 CPU 配置-1，比如：任务配置为 12C 时，建议该参数设置为 11 。另外，num_workers 数值可以适当调大，因为
dataset iter中有部分时间是在网络或者磁盘 IO， 这部分不消耗 CPU；但是也不能设置太大，因为数据预处理部分是 CPU 密集型任务，并行进程过多，会造成 CPU 争抢从而降低预处理效率。
优化案例一
线上一个基于 MMDetection 框架（其底层也是调用 PyTorch 框架）的 CV 模型训练任务，在做参数调整之前，单个 step 耗时不稳定，平均在 1.12s 左右，其中拉取数据时长在 0.3s 左右：
调整参数之后，单个 step 耗时稳定，平均在 0.78 s 左右，其中拉取数据耗时 0.004s，基本可以忽略。
该模型训练任务，通过上述优化调整，数据拉取时间缩短为 0，单个 step 的耗时从原来的 1.12s 降到 0.78s，整体训练时间减少 30%（从 2 天缩短到 33 小时），效果显著。
优化案例二
线上某个多模态模型（输入包含图片和文字）训练任务，使用 2 卡 V100 训练，参数调整如下：
调整后训练 300 step 总消耗时 405s，整体训练时间减少 45%左右（从 10 天缩短到 5 天左右）。
优化案例三
线上某 YoloX 模型训练任务，使用单卡 A100 训练，参数调整如下：
调整后整体训练时长减少 80%左右（从 10 天 19 小时，缩短至 1 天 16 小时）。
3.2.3 数据拉取 IO 瓶颈分析
当前，KubeAI 平台为训练场景提供 3 种存储介质：
本地盘：空间小，读写性能最好，单盘 500G~3T 空间可用。
NAS 网络存储：空间大，读写性能较差，成本适中。
CPFS 并行文件系统存储：空间大，读写性能好，成本高。
对于小数据集，可以先将数据一次性拉取到本地盘，然后每个 epoch 从本地盘来读数据，这样避免了每一个 epoch 重复的从远程 NAS 来拉取数据，相当于整个训练只需要从远程 NAS 拉取一次数据。对于大数据集，有 2 种解决方案：
将大数据集提前进行 resize，存储比较小的图片来进行训练，这样避免了每个 epoch 都需要 resize，而且 resize 之后，图片变小，读取更快。
将数据集放入并行文件系统 CPFS 存储上，提高训练吞吐。实验表明 CPFS 在图片场景下是 NAS 盘读性能的 3～6 倍。
3.3 TrainingModel 优化
数据部分优化后，训练过程中的主要时间开销就在 GPU 训练部分了。目前业内有一些比较成熟的方法可以参考，我们总结如下。
3.3.1 混合精度训练（AMP）
PyTorch 混合精度训练在 PyTorch 官网有详细介绍，以及开启混合精度训练的方法，可以阅读这里获取实现方法。当前许多 CV 训练框架已经支持 AMP 训练，比如：
MMCV 框架中 AMP 参数就是开启混合精度训练的选项。
Pytorch Vision 中也有相关参数来开启 AMP 训练。
需要说明的是，混合精度训练过程中并不是将所有模型参数都转为 FP16 来计算，只有部分做转换。混合精度之所以能加速训练过程，是因为大部分英伟达 GPU 机型在 FP16 这种数据格式的浮点算力比 FP32 要快一倍；此外，混合精度训练显存占用会更小。
3.3.2 单机多卡数据并行训练
Pytorch 原生支持多卡数据并行训练，详细开启多卡训练的方式参考官方文档。多卡训练过程中每一张卡的 backword 计算会多加一次多卡之间集合通讯 all-reduce 操作，用来计算多张卡上的梯度的平均值。
3.4 自研训练引擎框架 kubeai-training-framework
通过前面的分析我们可以看到，虽然 PyTorch 框架本身已经做的很好了，训练方式、参数支持丰富，但在实际的模型研究和生产过程中，由于模型的差异性、训练数据的差异性，以及模型开发者的经验差异性，PyTorch 框架本身的优势不一定能够发挥出来。
基于前述分析和实践，KubeAI 平台开发了训练引擎框架 kubeai-training-framework，帮助模型开发者更好地匹配训练脚本参数，快速接入使用合适的训练方式。_kubeai-training-framework_中包含 PyTorch Dataloader 优化、GPU TrainModel（AMP）提速以及各种功能函数等。以 Dataloader 为例，用户可通过以下方式使用：
4.AI Pipeline 引擎助力 AI 业务快速迭代
通常模型的开发可以归纳为如下图所示的过程：
可以看到，在需求场景确定、第一个模型版本上线之后，模型是需要反复迭代的，以期望取得更好的业务效果。KubeAI 平台在迭代建设的过程中，逐步上线了 Notebook、模型管理、训练任务管理、推理服务管理等一个个相对独立的功能模块。随着业务需求的不断变化，模型迭代效率直接影响了业务的上线效率，KubeAI 平台建设了 AI Pipeline 能力，重点解决 AI 场景的周期性迭代类需求，提高生产效率。
AI Pipeline 是在 ArgoWorkflow 基础上做了二次开发，以满足模型迭代、推理任务管理、数据处理等对定时需求、任务启动触发方式、通用模板任务、指定节点启动等需求。AI Pipeline 上线之前，一个迭代任务可能会被配置为多个分散的任务，维护工作量大，调试周期长。如下图是做一个类似任务需要单独配置的任务情况：
AI Pipeline 可以将整个工作流设计成如下图所示：
Pipeline 编排的方式，减少了模型开发者浪费在重复工作上的时间，可以将更多的时间投入到模型研究上。同时，通过合理编排任务，可以对有限的资源进行充分地利用。
5，展望
KubeAI 平台从得物 AI 业务场景的实际需求出发，以三大核心引擎为建设目标，着力解决 AI 模型研发过程中的训练、推理性能问题，以及模型版本迭代过程中的效率问题。
在推理服务性能上，我们会以 kubeai-inference-framework 为起点，继续在模型量化、算子优化、图优化等方面进行深入探索。在模型训练方面，我们会继续在图像数据预处理、Tensorflow GPU 训练框架支持、NLP 模型训练支持上发力，以 kubeai-training-framework 训练引擎框架为接口，为模型开发者提供更高效、性能更高的训练框架。此外，AI Pipeline 引擎上，我们会支持更丰富的预置模型，以满足通用数据处理任务、推理任务等需求。
文：伟东
本文属得物技术原创，来源于：得物技术官网
未经得物技术许可严禁转载，否则依法追究法律责任！ 著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
版权声明: 本文为 InfoQ 作者【得物技术】的原创文章。
原文链接:【http://xie.infoq.cn/article/6d6361558a41676456c53544e】。
本文遵守【CC BY-NC】协议，转载请保留原文出处及本版权声明。
得物APP技术部 2019-11-13 加入
关注微信公众号「得物技术」
评论