硬件平台上深度学习自动内核优化

硬件平台上深度学习自动内核优化

对于AI开发人员来说,在各种硬件平台上优化深度神经网络的性能仍然是一个难题。在系统支持方面,在这里面临着许多问题:将训练有素的模型从多个前端(例如Tensorflow,ONNX,MXNet)部署到多个硬件平台(例如CPU,GPU,加速器)。此问题最关键的性能部分,为不断增长的模型体系结构和硬件平台获得高性能的内核实现。

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

本文展示了TVM编译器堆栈中自动内核优化的工作流程,以及几种硬件平台上的基准测试结果。

系统总览

 硬件平台上深度学习自动内核优化

 

 图1.系统概述

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

内部循环使用可伸缩的RPC运行时,基于机器学习的调谐器和张量编译器。在循环的每一轮中,调谐器都会从较大的搜索空间中,选择一批有前途的候选内核实现,将其配置在实际硬件上。然后,调谐器将获得性能分析结果。这些分析结果用作训练数据,以适合预测模型。拟合预测模型后,调谐器根据预测,选择下一个有前途的候选者,然后循环继续进行。这样,迭代地快速搜索内核。

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

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

硬件平台上深度学习自动内核优化

 

 图2.传统自动调整和AutoTVM的比较

开始调整

为了演示,在ARM开发板RK3399上对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卷积层包含 13,000行代码。项目存储库中引用基准编号。如下图所示,TVM在Rasbperry Pi 3B上的所有网络上均胜过它。

 硬件平台上深度学习自动内核优化

 

 Mali马里GPU

ARM Compute Library是供应商提供的库,它很好地支持Mali GPU(OpenCL)。由于卷积层的优势,TVM在ResNet和MobileNet中提供了更强大的性能。TVM在vgg-16上落后一些,因为vgg-16是一个古老而庞大的网络,并且具有多个大型密集层。

 硬件平台上深度学习自动内核优化

 

 NVIDIA GPU

在NVIDIA GPU上,CuDNNTensorRT分别是两个供应商提供的用于训练和推理的库。由于专注于推理,因此在无限制的环境中运行基准测试。另一个张量编译器PlaidML也被报告为基线,因为与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] benchmark: https://github.com/dmlc/tvm/tree/master/apps/benchmark
[2] Tutorial on tuning for ARM CPU: https://tvm.apache.org/docs//tutorials/autotvm/tune_nnvm_arm.html
[3] Tutorial on tuning for Mobile GPU: https://tvm.apache.org/docs//tutorials/autotvm/tune_nnvm_mobile_gpu.html
[4] Tutorial on tuning for NVIDIA/AMD GPU: https://tvm.apache.org/docs//tutorials/autotvm/tune_nnvm_cuda.html
[5] Paper about AutoTVM: Learning to Optimize Tensor Program
[6] Paper about Intel CPU (by AWS contributors) : Optimizing CNN Model Inference on CPUs

上一篇:linux远程scp,但是不覆盖已经存在文件的方法rsync


下一篇:资源下载