CUDA上深度学习模型量化的自动化优化

CUDA上深度学习模型量化的自动化优化

深度学习已成功应用于各种任务。在诸如自动驾驶汽车推理之类的实时场景中，模型的推理速度至关重要。网络量化是加速深度学习模型的有效方法。在量化模型中，数据和模型参数都用诸如int8和float16低精度数据类型表示。降低的数据带宽减少了推理时间和存储器/存储要求，以及功耗。在适当的量化方案下，可以最小化量化模型的精度下降。因此，量化模型特别适合研究人员和开发人员，使大型模型适合在各种设备（例如GPU，CPU和移动设备）上部署。

通常通过手工微内核，针对不同的工作量来优化量化算子，或者依赖于cuDNN和TensorRT等黑盒专有解决方案。在集成中编写高性能微内核，可能非常困难，通常需要大量的工程工作。此外，很难使这些临时微内核，适应新兴的工作负载和新设备。

图1. TVM，TensorRT和MXNet上不同模型的推理时间

TVM通过全栈编译器和基于机器学习的优化器，自动生成计算内核，从而解决了这一难题。TVM可以通过在人为设计的搜索空间，进行自动搜索来生成有效的内核。在VGG和ResNet等标准工作负载中，TVM与其它最新框架相比具有竞争优势。在诸如ResNeXt和Deformable ConvNets之类的新兴模型中，自动优化使TVM可以轻松适应这些新工作负载并显着提高性能。

本文将展示如何使用TVM在CUDA上自动优化量化深度学习模型。

在TVM中表达量化的CUDA内核

通过张量化利用张量本征

许多平台都为特殊的计算模式，提供了特定于体系结构的指令，例如x86上的SIMD指令dp4a以及hfmaCUDA上的指令。这些内部指令针对特定设备进行了高度优化。通过利用硬件内部函数，可以为量化算子带来显着的性能提升。

目前，dp4a已被CUDA上的TVM int8 算子广泛使用。dp4a是Compute Capability 6.1设备上固有的CUDA。这是一种混合精度指令，可以高效地计算两个4元素8位整数向量之间的点积，并以32位格式累加结果。使用dp4a，可以在8位整数向量之间实现点积，其元素数可以被4整除。使用高效的点积算子，可以实现高级算子，例如2D卷积和密集层，这些算子通常由点积支持。

为了说明这一点，在二维卷积中，沿着内核的通道，宽度和高度进行累积。这是的典型用例dp4a。TVM使用张量化来支持调用外部内部函数。不需要修改原始的计算声明；使用调度原语tensorize，用dp4a张量内部函数代替累加。

数据布局重排

张量化的挑战之一，可能需要设计特殊的计算逻辑，适应张量本征的要求。尽管在密集算子中沿张量的内轴累积是很自然的，conv2d但可能更具挑战性。在conv2d期望采取切片作为输入端的通道尺寸dp4a，因为信道数目通常是4的倍数（否则回落到原来conv2d在NCHW布局）。为了实现内存局部性，先沿最内轴减小。考虑到这些因素，使用自定义数据布局来应对这一挑战。

在CUDA int8 2d卷积中，根据经验选择NCHW4c作为数据布局和OIHW4o4i权重布局。模板可以很容易地推广到NCHW[x]c和OIHW[x]o[x]i，其中x是可以被4整除的任意正整数。选择数据布局中，通道切片位于打包的最里面维度中。同样，在权重的输入和输出通道尺寸中打包切片，以使输出与输入具有一致的数据布局，从而防止了图层之间的多余布局转换。

在图2中显示了2d卷积输出的一个元素的计算。超尺寸（包含打包元素的块状布局的外部尺寸）每个位置的元素NCHW和OIHW分别是打包的输入和内核。打包内核的每一列都来自不同的过滤器。使用dp4a来计算压缩输入和压缩内核中每一行之间的点积，将结果累加到输出张量。

图2. NCHW4c中的数据布局和OIHW4o4i中的权重布局的二维卷积。 左：NCHW4c布局中的输入张量。内核的一个移动的过滤器被染成蓝色。输入和内核的一个元素以灰色显示。 中：灰色块中的打包输入和内核。 右：NCHW4c布局中的输出。所示一个元素内，通道子维度中有四个压缩元素。

在指定卷积层的布局之后，其它算子（例如add和激活），可以在Relay的AlterOpLayout传递期间，自动适应所选的布局。权重的布局转换，可以离线进行预先计算。可以在相同的布局中，运行整个模型而无需额外的开销。

设计搜索空间以进行自动优化

在量化算子中获得良好性能的关键是与基于机器学习的自动优化相集成。一个问题是如何设计有效的调度搜索空间。有效的调度模板意味着在自动调整中，可以通过合理的迭代次数获得良好的性能。一般而言，努力定义一个灵活的模板，覆盖搜索空间中的不同配置。另一方面，利用性能优化方面的先验知识。例如，由于在CUDA编程中将数据缓存在共享内存中是一种常见的做法，利用共享内存，使用机器学习来选择最佳的图块大小。还进行了一些手动平铺，例如将轴分割为4或16，以便于向量化内存访问。

在量化的2d卷积中，设计了一个搜索空间，其中包括一组可调选项，例如图块大小，要融合的轴，循环展开和双缓冲的配置。在CUDA上注册了量化模板conv2d和denseCUDA模板int8。在自动调整期间，可以通过设置template_key参数来为这些量化算子创建调整任务。

一般工作流程

图3.运行量化模型的工作流程

TVM提供了一个简单的工作流程，可以从其它框架量化训练后的模型，自动优化算子（使用AutoTVM），部署到其它设备。

首先，使用Relay前端导入现有模型。在这里，以带有(1, 3, 224, 224)输入形状的MXNet模型为例。

sym, arg_params, aux_params = mxnet.model.load_checkpoint(model_path, epoch)

net, params = relay.from_mxnet(sym, shape={'data': (1, 3, 224, 224)}, arg_params=arg_params, aux_params=aux_params)

接下来，使用Relay量化API将其转换为量化模型。

net = relay.quantize.quantize(net, params=params)

然后，使用AutoTVM为模型中的算子，提取调整任务并执行自动优化。

最后，建立模型并以量化模式运行推理。

with relay.build_config(opt_level=3):

graph, lib, params = relay.build(net, target)

结果relay.build是可部署的库。可以直接在GPU上运行推理，也可以通过RPC部署在远程设备上。

基准

为了验证TVM中量化算子的性能，对几种流行的网络模型（包括VGG-19，ResNet-50和Inception V3）的性能进行了基准测试。还对来自可变形卷积网络的DRN-C-26，ResNeXt-50和DCN-ResNet-101进行基准测试，以显示新兴模型的性能，其中包含较少的常规算子，例如膨胀卷积，组卷积和可变形卷积。选择NVIDIA TensorRT作为基准。据报道，在float32模式下MXNet 1.4 + cuDNN 7.3的结果显示了量化的加快。实验是在NVIDIA GTX 1080上进行的。当以批处理大小=1和16运行时，报告每个图像的推理时间。

如图1所示，TVM使用量化实现了高达8倍的加速。在VNG和ResNet等标准CNN模型中，TVM与TensorRT的最新结果保持一致。

在对新兴模型进行基准测试时，TVM取得了令人印象深刻的结果。在ResNeXt和DCN-ResNet-101上获得了可观的性能提升。TensorRT的DCN-ResNet-101的结果不可用，因为尚无可变形卷积的正式实现。TVM中的自动优化，可以轻松灵活地支持和优化新兴的工作负载。