[论文解读] Learning to Optimize Tensor Programs

Tianqi Chen, Lianmin Zheng+6 more|arXiv (Cornell University)|May 21, 2018

Parallel Computing and Optimization Techniques参考文献 36被引用 124

一句话总结

本文提出 AutoTVM，一种基于学习的框架，构建领域特定的成本模型来自动搜索并优化深度学习工作负载的张量程序实现，通过迁移学习在 CPU、GPU 以及移动硬件上实现具有竞争力的性能。

ABSTRACT

We introduce a learning-based framework to optimize tensor programs for deep learning workloads. Efficient implementations of tensor operators, such as matrix multiplication and high dimensional convolution, are key enablers of effective deep learning systems. However, existing systems rely on manually optimized libraries such as cuDNN where only a narrow range of server class GPUs are well-supported. The reliance on hardware-specific operator libraries limits the applicability of high-level graph optimizations and incurs significant engineering costs when deploying to new hardware targets. We use learning to remove this engineering burden. We learn domain-specific statistical cost models to guide the search of tensor operator implementations over billions of possible program variants. We further accelerate the search by effective model transfer across workloads. Experimental results show that our framework delivers performance competitive with state-of-the-art hand-tuned libraries for low-power CPU, mobile GPU, and server-class GPU.

研究动机与目标

• 促进超越手工调优库的张量算子自动优化。
• 形式化在一个庞大张量程序计划空间中探索以最小化硬件运行时间的问题。
• 开发基于学习的成本模型以引导对数十亿种程序变体的搜索。
• 实现迁移学习，以在不同工作负载和硬件目标之间重用知识。

提出的方法

• 为给定的指表达式 e 和代码生成器 g 定义调度搜索空间 Se，成本 f(x) 未知但可测量。
• 开发两种成本模型：使用来自低级 AST 的领域特征的梯度提升树（GBT），以及对 AST 的 TreeGRU 嵌入。
• 用回归或排序为基础的目标来训练模型以预测运行时间并引导选择。
• 使用带有模拟退火和多样性感知选择的探索循环来生成候选调度并在硬件上测量它们的实际性能。
• 通过将全局不变表示模型与局部领域内模型相结合来实现迁移学习，从而加速新工作负载的优化。
• 报告 AutoTVM（TVM 基于）的端到端增益在各后端实现，并且无需外部算子库。

实验结果

研究问题

• RQ1 基于学习的成本模型是否能够有效地在十亿级规模的计划空间中导航，以在真实硬件上找到快速的张量程序实现？
• RQ2 迁移学习是否能够在不同工作负载和硬件目标之间实现更快的优化？
• RQ3 在该领域中，回归与排序中哪种目标更有利于发现快速的调度？
• RQ4 在这种设置下，多样性感知的探索和不确定性估计是否有帮助？

主要发现

• 统计成本模型（GBT 和 TreeGRU）在以更少的硬件评估找到更快张量程序方面优于黑箱基线。
• 基于排序的目标在按速度排序候选程序方面通常与回归目标同样好或更好。
• 多样性感知探索的影响几乎为零到中等，在某些工作负载上有一定提升。
• 迁移学习通过在 GPU 和移动设备之间重复使用先前看到的工作负载知识，实现 2x 到 10x 的加速。
• AutoTVM 的端到端推理提升在 NVIDIA TITAN X、ARM Cortex-A53 和 ARM Mali 后端的范围为 1.2x 到 3.8x。