面向所有硬件平台的深度学习自动内核优化

2018年10月3日 • Lianmin Zheng, Eddie Yan, Tianqi Chen

对于 AI 开发者来说，在各种硬件平台上优化深度神经网络的性能仍然是一个难题。在系统支持方面，我们面临着一个多对多的问题：将来自多个前端（例如 Tensorflow、ONNX、MXNet）的训练模型部署到多个硬件平台（例如 CPU、GPU、加速器）。此问题中最关键的性能部分是为不断增长的模型架构和硬件平台获得高性能内核实现。

为了应对这一挑战，TVM 采用了全栈编译器方法。TVM 结合了代码生成和自动程序优化，以生成与经过大量手工优化的库相媲美的内核，从而在包括 ARM CPU、Intel CPU、Mali GPU、NVIIDA GPU 和 AMD GPU 在内的硬件平台上获得最先进的推理性能。

在这篇博客文章中，我们将展示 TVM 编译器堆栈中自动内核优化的工作流程，以及在多个硬件平台上的基准测试结果。

系统概述

TVM 中的内核优化以迭代循环方式完成。如图 1 所示，自动内核优化以来自前端框架的神经网络（通常以计算图表示）作为输入，并为该网络中的所有运算符生成内核。

内部循环使用可扩展的 RPC 运行时、基于机器学习的调优器和张量编译器。在循环的每一轮中，调优器从大型搜索空间中选择一批有希望的候选内核实现，并在真实硬件上对其进行性能分析。然后，调优器获得性能分析结果。这些性能分析结果用作训练数据，以拟合预测模型。在拟合预测模型后，调优器根据预测选择下一个有希望的候选者，循环继续。这样，我们迭代地搜索快速内核。

下图比较了传统的自动调优和 AutoTVM。主要区别在于 AutoTVM 是

• 可扩展的，适用于异构设备集群
• 学习使用可迁移的机器学习成本模型优化张量程序

您可以参考我们的论文 [1] 了解更多详情。

开始调优

为了演示，我们在 RK3399（一个 ARM 开发板）上运行 resnet-18 的优化。由于博客文章篇幅有限，此处省略了详细说明。ARM CPU、Mali GPU、NVIDIA GPU、AMD GPU 的教程链接都在本博客的末尾提供。

首先，我们从 MXNet 模型库中获取一个预训练模型，并从中提取调优任务。

from mxnet.gluon.model_zoo.vision import get_model

block = get_model('resnet18_v1', pretrained=True)
net, params = nnvm.frontend.from_mxnet(block)

tasks = autotvm.extract_from_graph(net)
tune_tasks(tasks, **tuning_option)

resnet-18 中有 12 个不同的 conv2d 层，因此我们启动 12 个调优任务。对于每个任务，调优器进行数百次试验并选择最佳的一个。完成所有调优任务后，我们编译整个网络并生成一个可部署的最小库。一个示例输出是

Extract tasks...
Tuning...
[Task  1/12]  Current/Best:   22.37/  52.19 GFLOPS | Progress: (544/1000) | 406.59 s Done.
[Task  2/12]  Current/Best:    6.51/  18.77 GFLOPS | Progress: (608/1000) | 325.05 s Done.
[Task  3/12]  Current/Best:    4.67/  24.87 GFLOPS | Progress: (480/1000) | 372.31 s Done.
[Task  4/12]  Current/Best:   11.35/  46.83 GFLOPS | Progress: (736/1000) | 602.39 s Done.
[Task  5/12]  Current/Best:    1.01/  19.80 GFLOPS | Progress: (448/1000) | 262.16 s Done.
[Task  6/12]  Current/Best:    2.47/  23.76 GFLOPS | Progress: (672/1000) | 563.85 s Done.
[Task  7/12]  Current/Best:   14.57/  33.97 GFLOPS | Progress: (544/1000) | 465.15 s Done.
[Task  8/12]  Current/Best:    1.13/  17.65 GFLOPS | Progress: (576/1000) | 365.08 s Done.
[Task  9/12]  Current/Best:   14.45/  22.66 GFLOPS | Progress: (928/1000) | 724.25 s Done.
[Task 10/12]  Current/Best:    3.22/  15.36 GFLOPS | Progress: (864/1000) | 564.27 s Done.
[Task 11/12]  Current/Best:   11.03/  32.23 GFLOPS | Progress: (736/1000) | 635.15 s Done.
[Task 12/12]  Current/Best:    8.00/  21.65 GFLOPS | Progress: (1000/1000) | 1111.81 s Done.
Compile...
Upload...
Evaluate inference time cost...
Mean inference time (std dev): 162.59 ms (0.06 ms)

如果您的模型具有一些奇怪的形状或您的硬件是定制的，则调优尤其有帮助且值得一试，因为手工优化的静态库无法考虑所有情况。

基准测试结果

我们在设备集群上预先调优了一些流行的网络，并发布了以下基准测试。重现说明在本博客的末尾。

由于我们有一个统一的运行时接口，因此全面基准测试 TVM 很简单。然而，如果没有许多其他项目的开发人员的专家协助，维护与其他所有平台的完整、最新和正确的比较是不可行的。因此，我们将所有数字放在一个表格中，然后提供与其他一些库的不完整比较。

比较

我们通过在每个平台上与经过大量优化的传统库进行比较，验证了我们自动优化堆栈的有效性。

我们在 ImageNet (3x224x224) 数据集上测试了流行的图像分类网络，批量大小为 1，数据类型为 float32。报告的数字是每张图像的时间成本，以毫秒为单位。

ARM CPU

我们选择 NCNN，一个广泛使用的、手工优化的内核库作为基线。它广泛使用了 NEON 汇编指令。例如，代码库仅针对 3x3 卷积层就包含 1.3 万行代码。我们参考了他们项目存储库中的基准测试数字。如下图所示，TVM 在 Rasbperry Pi 3B 上针对所有网络都优于它。

Mali GPU

ARM Compute Library 是供应商提供的库，可以很好地支持 Mali GPU (OpenCL)。根据结果​​，由于卷积层的优势，TVM 在 ResNet 和 MobileNet 中提供了更强大的性能。TVM 在 vgg-16 上稍有落后，因为 vgg-16 是一个旧的且庞大的网络，并且有几个大型密集层。

NVIDIA GPU

在 NVIDIA GPU 上，CuDNN 和 TensorRT 是两个供应商提供的分别用于训练和推理的库。由于我们专注于推理，因此我们在非批量设置中运行基准测试。另一个张量编译器 PlaidML 也被报告为基线，因为之前有将其与 pre-AutoTVM 版本的 TVM 进行比较的基准测试。我们从 PlaidBench 中参考了其基准测试结果。根据以下结果，TVM 实现了与 TensorRT 性能相当的水平。

AMD GPU

我们还快速了解了一下 AMD GPU。TVM 支持 OpenCL 和 ROCm 后端。我们发现 ROCm 更好，因为它更专门用于 AMD GPU。MIOpen 是供应商提供的内核库。TVM 的图形运行时可以直接调用 MIOpen 的内核实现，因此我们通过使用此集成来报告基线性能。

我们没有对 AMD GPU 进行任何特定的优化。NVIDIA GPU 的所有计算定义和调度代码都直接重用。因此，在大多数情况下，TVM 比 MIOpen 稍慢。我们相信仍有改进的空间。

我们的所有结果

我们在 ImageNet (3x224x224) 数据集上测试了以下网络，批量大小为 1，数据类型为 float32。报告的数字是每张图像的时间成本，以毫秒为单位。

• 注 1：此板内存不足。
• 注 2：由于时间限制，我们没有在 GPU 上调优一些小型网络。当性能分析数据不可用时，TVM 可以使用回退代码生成。但在此情景下，不保证具有竞争力的性能。

结论

凭借富有表现力的代码生成器和高效的搜索算法，我们能够生成与经过大量手工优化的内核相媲美的内核。由于程序员的时间很昂贵，而机器时间越来越便宜，我们相信在循环中使用真实硬件和数据的自动优化将成为推理部署的标准工作流程。TVM 正是提供了这样的解决方案。

链接

[1] 基准测试：https://github.com/dmlc/tvm/tree/master/apps/benchmark
[2] ARM CPU 调优教程：https://tvm.apache.org/docs//tutorials/autotvm/tune_nnvm_arm.html
[3] 移动 GPU 调优教程：https://tvm.apache.org/docs//tutorials/autotvm/tune_nnvm_mobile_gpu.html
[4] NVIDIA/AMD GPU 调优教程：https://tvm.apache.org/docs//tutorials/autotvm/tune_nnvm_cuda.html
[5] 关于 AutoTVM 的论文：Learning to Optimize Tensor Program
[6] 关于 Intel CPU 的论文（由 AWS 贡献者撰写）：Optimizing CNN Model Inference on CPUs