背景

CV 模型是业务中常见的模型，但是我们观察到 UC 集群中的 CV 类模型还有很大的 GPU 利用率提升空间。如果不对此进行优化，则需要大量的 GPU 资源才能才能满足延迟要求。为了提高资源利用率、节省计算资源，我们对这部分模型（主要是 PyTorch）的优化进行了一些探索，并针对使用资源较多的 Detectron2 模型进行测试。

Detectron2采用centernet作为backbone 提供特征提取能力，提取的特征经过proposal_generator 生成候选框并对候选框进行调整，排重和初步过滤，最后由 roi_heads 进行打分排序输出目标检测的score 和boxes 。

模型计算耗时主要集中在backbone部分。线上采用Python直接部署，优化前单卡服务能力较弱，单卡超过12qps的时候开始出现排队情况，并且在qps 20 后 延迟急剧放大甚至服务异常。为了维持服务稳定运行，线上只能使用较多的计算资源来提升qps，请求高峰期时，服务的qps在1850左右。 但这种以增加资源的方式来提升qps的做法会造成大量的资源浪费，经观察，在高峰期gpu的使用率只能维持在20%左右。

优化方案探索

优化主要考虑两个方向：1）提升单 CUDA Stream 条件下模型的推理速度；2）利用多CUDA Stream进一步提升GPU利用率。

模型推理后端优化

优化后端主要关注点

作为工程团队，我们比较关注优化工具的如下特性：

• 优化的成功率：PyTorch具有灵活、扩展性强、OP种类丰富等特点，这会导致其建模、推理代码的多样性很高。而另一方面，多数后端支持的算子类型仅为PyTorch的一个子集。如何在这种情况化最大化优化的成功率是最需要考量的一点。
• 优化的性能：一方面，可以充分发挥部署硬件的能力，降低用户响应时间，同时节省成本；另一方面，可以在保持响应时间不变的前提下，使用结构更为复杂的深度学习模型，进而提升业务精度指标。
• 优化及部署难易程度：优化工具需要提供简洁、稳定、易于使用的API，能够使得其使用者快速地完成模型优化的流程。同时，优化后的模型需要能够较容易地部署至现有工程体系中。

几类主流的优化手段

原生TensorRT

从优化的性能角度而言，CV领域模型当前的最优做法是全图转TensorRT，以充分利用其加速特性，团队内也已沉淀了许多TensorRT的加速经验并累积了一些TensorRT插件。要将PyTorch模型转化至TensorRT，有两种主流的方式：

1、PyTorch模型导出ONNX，再由onnx2trt将其解析成TensorRT的网络；

2、直接调用TensorRT network的API构建其等价模型，并加载参数。

上述方式都存在的问题是：TensorRT支持的算子种类的少于ONNX，同时PyTorch与ONNX之间也存在许多算子无法转化。因此这两种方式的优化成功率会较低。此外，TensorRT某些Pluginin的实现还存在鲁棒性问题，在不同模型上可能会存在数值差异（比如roiAlign）。

TensorRT-based Backend

考虑到PyTorch算子的多样性，我们进一步尝试基于子图架构的TensorRT-based backend。子图是指在优化前的计算图中进行可优化OP的选择，同时将这些聚合成子图，每个子图会被转化至相应的后端进行优化。而无法转化至相应后端的op，则回退至PyTorch执行。

这种架构可以在保证部署成功率的前提下获得最大的优化效果。目前社区中主流的以圈图架构接入TensorRT的优化工具包括两款，即Torch-TensorRT以及PAI-Blade。

图转化架构上而言，Torch-TensorRT支持TorchScript以及Torch Fx Module两种输入。它会针对PyTorch OP编写许多Converter，在Converter中调用TensorRT netowrk的API进行OP翻译并构建网络。PAI-Blade支持TorchScipt作为输入，对于子图TorchScript会经过ONNX并转化为TensorRT network。经由ONNX的路线可以复用社区成熟的工具，但多一层中间表示会使得定制化开发的链路较长；直接写converter可定制化链路缩短，但开发量大大提升，需要针对每个OP以及其不同的状况从头开发。当前阶段而言，经由ONNX路线的成熟度是更高的。

TensorRT功能覆盖度而言，二者皆提供了dynamic range、低精度优化、自定义plugin等功能。然而，Torch-Tensor如果需要启用dynamic range、int8量化等高阶功能，需要输入的模型能够做到全网转化到TensorRT，而PAI-Blade则在支持在圈图的架构下，每个子图单独启用这些高阶功能。

优化成功率以及优化性能方面，PAI-Blade针对TorchScript的特性开发了大量的JIT passes使得其优化后的图更为干净高效，圈图的范围也更大，进而可以获得更大的优化效果。但是由于ONNX这一层中间媒介的存在，使得一些冷门算子无法被转化至TensorRT后端。而Torch-TensorRT当前阶段的成熟度较低，优化成功率也较低。

除去TensorRT外，我们还调研了OneFlow、PaddlePaddle等业界比较流行的后端，这些后端都存在着和原生TensorRT 类似的问题，如算子支持不完整、opset 比较低等问题。

经过一段时间的对比，综合各方面考虑，最终选择易用性、鲁棒性和性能更好的 PAI-Blade 优化方案， 并对原始模型实现进行重写和改图适配。

多 CUDA Stream 优化

PyTorch 模型将所有线程的cuda 操作放到default stream， 导致多线程跑的时候 在CUDA kernel执行上串行化，而且还阻塞其他stream的执行，多线程内部执行的推理都是相互独立的个体，执行没有上下有依赖，因此可以对stream打散处理。一个典型的单CUDA Stream的timeline如下：

优化过程

模型推理代码重写

PyTorch模型在nn.Module层面是不包含一张“显式”的计算图的，需要将其转化成为TorchScript以获得推理计算图，进而进行推理优化。一方面，TorchScript仅为Python的一个子集，支持的语法受限。另一方面，PyTorch动态图的特性会导致模型的推理代码多样性非常大。这个冲突会导致在不修改推理代码的前提下，无法导出TorchScript。因此我们需要对Detectron2推理代码做出修改，主要包含如下几个部分：

语法和逻辑重写

      1. 变量、参数显式增加类型标注

            torchscript 会将没有显式标注的参数类型当做Tensor，如果实际变量未标注，同时类型不为Tensor的话，推理的过程出现类型不匹配的异常。

# 标注前
def forward(self, x, dim):
    return torch.sum(x, dim=dim)
# 标注后
def forward(self, x: torch.Tensor, dim: int):
    return torch.sum(x, dim=dim)

      2. 诸如Dict、List等容器类型的存储内容保持一致

TorchScript要求容器对象内存储的数据类型完全一致，比如一个List中所有数据的类型需要保持一致。

# 修改前
x = [1, "str"]
num_data = x[0]
str_data = x[1]
# 修改后
int_x = [1,]
str_x = ["str",]
num_data = int_x[0]
str_data = str_x[0]

     3.  对于__getattr__等TorchScript不支持的magic method进行重写

def init():
      self.local_modules :nn.ModuleList = []
      self.local_modules_tmp :Dict[str, nn.Module] = {}
      #self.add_module("acb", lateral_conv)
      self.local_modules_tmp["acb"] = lateral_conv.cuda()
      self.local_modules = nn.ModuleDict(self.local_modules_tmp)
def forward(self, input_node):
     for kn, vv in self.local_modules.items():
         if name == kn:
              input_node = vv(input_node)
     #input_node = self.__getattr__(name)(input_node)

     4. * 展开功能替换

boxlists = list(zip(*sampled_boxes)) # 替换为for循环

     5. list、dict、NoneType 等类型修改。list替换为 torch.nn.ModuleList，dict 替换为 torch.nn.ModuleDict等。

     6. 善用bypass的机制来忽略推理无关的代码分支

模型的推理代码可能同时存在着训练和推理相关的部分，同时还会包括前后处理等不属于模型推理本身的内容。这部分代码可以直接通过相关机制进行bypass。

# 修改前
if self.traininig:
    do_some_training_thing
else:
    do_some_eval_thing
# 修改后
if self.traininig:
    assert not torch.jit.is_scripting(), "Bypass the codes blow"
    do_some_training_thing  # assert之后的代码均不会被翻译
else:
    do_some_eval_thing

     7.  目前部署环境目前都用单卡推理，可以使用 torch.nn.BatchNorm2d 代替 torch.nn.SyncBatchNorm。

模型参数的key修改

  在模型代码修改完毕后，由于torch.nn.Module的组织形式会发生一些变化，会使得其在state_dict中的key发生变化，由于存在key mismatch的情况，无法直接加载预训练模型。因此我们需要修改预训练模型的key以达到正确加载的目的。

self.linear = xxx  # key : “linear”
self.linear = torch.nn.ModuleList([xxx])  # key: "linear.0","linear.1",..

 拆图

实际测试中，并不是全图优化后的端到端性能是最好的，此外GPU运行的性能不一定由于CPU。因此，我们针对不同的网络组建(backbone, neck,  detection head)以及不同的代码逻辑进行了拆图，每个部分采取各自更优的方式进行优化。

转换&圈图分析

PAI-Blade优化

import torch_blade
import torch_blade.tensorrt
model_path = "./opt_model_backbone.pt"
# 加载优化前TorchScript
script_model = torch.jit.load(model_path)
script_model.cuda().eval()
...
# 配置优化
cfg = torch_blade.Config()
cfg.optimization_pipeline = "TensorRT" # 采取TensorRT后端
# 执行优化 
with cfg: 
  opt_model = torch_blade.optimize(script_model, False, model_inputs=img)
    print("Optimize finish!")
  # 优化后的模型同样是TorchScript，可以直接调用相关api序列化到本地
    torch.jit.save(opt_model, opt_file_name)
    print("save finish!")
# 加载模型并推理
opt_model = torch.jit.load(opt_file_name)
# 执行推理
opt_model(img)

 圈图分析

PAI-Blade圈图默认是最少三个基础的op，对于费时的op如果没有圈入到trt-grp执行，可以手工设置白名单，然圈图优化对该OP进行特殊处理，另外如果优化后性能不及预期，可以通过圈图信息分析查看是否存在关键子图未被圈入trt-grp

多CUDA Stream    

优化结果

两项优化在 Detectron2 模型上线之后，我们看到系统的 QPS 和 GPU 利用率均有了大幅提升。优化后可以节省接近 65% 的计算资源。

优化前

优化后

此外，我们也评估了 PAI-Blade 在其他典型模型上的优化效果：

随着优化在其他模型上的逐渐上线，预计将带来整体160+张GPU的优化收益。

总结

作为 UC 的模型推理优化团队，在不影响模型业务效果的情况下提高模型推理性能和减少计算资源使用是我们的主要目标。 在这次优化中，我们从优化的成功率、优化的性能、优化及部署难易程度三个角度对优化工具进行选型，最终选择了 PAI-Blade 集成到系统中。从最终的性能优化效果及资源成本的节省来看，基本达到了最初的预期。但性能优化没有终点，随着模型结构的不断进化、模型的规模逐渐变大，新的挑战还会不断出现，但优化工具的成功率、绝对性能、易用性仍然是我们技术选型的核心标准。