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

•


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

以前，量化算子通常使用手工制作的微内核针对不同的工作负载进行优化，或者依赖于 cuDNN 和 TensorRT 等黑盒专有解决方案。用汇编语言编写高性能微内核可能非常具有挑战性，并且通常需要大量的工程投入。此外，将这些临时的微内核适配到新兴的工作负载和新设备上也很困难。

image

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

TVM 通过全栈编译器和基于机器学习的优化器来解决这一挑战，以自动生成计算内核。TVM 可以通过在人工设计的搜索空间中自动搜索来生成高效的内核。在 VGG 和 ResNet 等标准工作负载中，TVM 实现了与其他最先进框架相比具有竞争力的性能。在新兴模型（如 ResNeXt 和 Deformable ConvNets）中，自动优化使 TVM 能够轻松适应这些新的工作负载并实现显着的性能提升。

在这篇文章中，我们将展示如何使用 TVM 在 CUDA 上自动优化量化深度学习模型。

在 TVM 中表达量化 CUDA 内核

通过张量化利用张量内在函数

许多平台为特殊的计算模式提供了特定于架构的指令，例如 x86 上的 SIMD 指令和 CUDA 上的 dp4a 和 hfma 指令。这些内在指令针对特定设备进行了高度优化。通过利用硬件内在函数，我们可以显着提高量化算子的性能。

目前，dp4a 已在 TVM CUDA int8 算子中得到广泛使用。dp4a 是 Compute Capability 6.1 设备上的 CUDA 内在函数。它是一种混合精度指令，可有效计算两个 4 元素 8 位整数向量之间的点积，并将结果累加为 32 位格式。使用 dp4a，我们可以实现元素数量可被 4 整除的 8 位整数向量之间的点积。有了高效的点积算子，我们可以实现诸如 2d 卷积和密集层等高层算子，因为这些算子通常由点积支持。

为了说明，在 2d 卷积中，我们沿着内核的通道、宽度和高度轴累加。这是 dp4a 的典型用例。TVM 使用张量化来支持调用外部内在函数。我们不需要修改原始计算声明；我们使用调度原语 tensorize 将累加替换为 dp4a 张量内在函数。有关张量化的更多详细信息，请参见教程。

数据布局重排

张量化中的挑战之一是我们可能需要设计特殊的计算逻辑来适应张量内在函数的要求。虽然在密集算子中沿着张量的内轴累加是很自然的，但 conv2d 可能更具挑战性。在 conv2d 中，我们希望在通道维度中取一个切片作为 dp4a 的输入，因为通道数通常是 4 的倍数（否则我们回退到 NCHW 布局中的原始 conv2d）。同时，为了实现内存局部性，我们希望首先沿最内轴进行归约。考虑到这些因素，我们使用自定义数据布局来应对这一挑战。

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

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

image

图 2. 数据布局为 NCHW4c，权重布局为 OIHW4o4i 的 2D 卷积。左：NCHW4c 布局中的输入张量。内核的一个移动滤波器以蓝色着色。输入和内核的一个元素以灰色着色。中：灰色块中的打包输入和内核。右：NCHW4c 布局中的输出。在描绘的一个元素内部，通道子维度中有四个打包元素。

在我们指定了卷积层的布局之后，add 和激活等其他算子可以在 Relay 的 AlterOpLayout 过程中自动适应所选布局。权重的布局转换可以离线预先计算。因此，我们可以在相同的布局中运行整个模型，而不会产生额外的开销。

为自动优化设计搜索空间

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

在量化 2d 卷积中，我们设计了一个搜索空间，其中包括一组可调选项，例如 tile 大小、要融合的轴、循环展开和双缓冲的配置。CUDA 上量化 conv2d 和 dense 的模板在模板键 int8 下注册。在自动调优期间，我们可以通过设置 template_key 参数为这些量化算子创建调优任务。有关如何启动自动优化的详细信息，请参见 AutoTVM 教程。

通用工作流程

image

图 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 提取模型中算子的调优任务并执行自动优化。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 在内的几个流行的网络模型进行了基准测试。我们还对来自 Deformable ConvNets 的 DRN-C-26、ResNeXt-50 和 DCN-ResNet-101 进行了基准测试，以展示新兴模型的性能，其中包含不太常规的算子，例如空洞卷积、分组卷积和可变形卷积。我们选择 NVIDIA TensorRT 作为我们的基线。报告了 MXNet 1.4 + cuDNN 7.3 在 float32 模式下的结果，以显示量化的加速效果。实验在 NVIDIA GTX 1080 上进行。我们报告了批量大小 = 1 和 16 时每个图像的推理时间。

如图 1 所示，TVM 使用量化实现了高达 8 倍的加速。在 VGG 和 ResNet 等标准 CNN 模型中，TVM 实现了与 TensorRT 最先进结果相当的性能。

在对新兴模型进行基准测试时，TVM 取得了令人印象深刻的结果。我们在 ResNeXt 和 DCN-ResNet-101 上获得了显着的性能提升。TensorRT 的 DCN-ResNet-101 结果不可用，因为没有可变形卷积的官方实现。我们表明，TVM 中的自动优化使其易于灵活地支持和优化新兴的工作负载。

代码示例

作者简介 & 致谢

吴伟林 是上海交通大学的本科生。他目前在 TuSimple 实习。作者非常感谢 陈天奇 和 Eddie Yan 的审阅。