[论文解读] Reinforcement Learning and Adaptive Sampling for Optimized DNN Compilation

Byung Hoon Ahn, Prannoy Pilligundla+1 more|arXiv (Cornell University)|May 30, 2019

Advanced Neural Network Applications参考文献 65被引用 13

一句话总结

该论文提出 ReLeASE，一种基于强化学习（RL）的深度神经网络（DNN）优化编译器，通过聚类实现自适应采样，以减少昂贵的硬件测量。通过将 DNN 编译建模为 RL 问题，并将采样聚焦于具有代表性的子空间，ReLeASE 在优化时间上比 AutoTVM 快 4.45 倍，推理性能最高提升 5.6%，同时显著减少搜索步数和测量开销。

ABSTRACT

Achieving faster execution with shorter compilation time can enable further diversity and innovation in neural networks. However, the current paradigm of executing neural networks either relies on hand-optimized libraries, traditional compilation heuristics, or very recently, simulated annealing and genetic algorithms. Our work takes a unique approach by formulating compiler optimizations for neural networks as a reinforcement learning problem, whose solution takes fewer steps to converge. This solution, dubbed ReLeASE, comes with a sampling algorithm that leverages clustering to focus the costly samples (real hardware measurements) on representative points, subsuming an entire subspace. Our adaptive sampling not only reduces the number of samples, but also improves the quality of samples for better exploration in shorter time. As such, experimentation with real hardware shows that reinforcement learning with adaptive sampling provides 4.45x speed up in optimization time over AutoTVM, while also improving inference time of the modern deep networks by 5.6%. Further experiments also confirm that our adaptive sampling can even improve AutoTVM's simulated annealing by 4.00x.

研究动机与目标

• 为减少当前依赖启发式或随机方法（如模拟退火和遗传算法）的 DNN 编译优化所耗费的时间与成本。
• 通过最小化所需采样次数，解决搜索过程中真实硬件测量带来的高计算成本问题。
• 通过采用强化学习而非传统搜索策略，提升 DNN 编译器优化的搜索效率与解的质量。
• 开发一种采样策略，在减少测量次数的同时保持高代表性与相关性，通过基于聚类的子空间聚焦实现。

提出的方法

• 将 DNN 编译建模为强化学习（RL）问题，其中智能体学习选择最优编译器参数（knobs）以最大化性能。
• 提出一种自适应采样算法，通过聚类搜索空间以识别具有代表性的子空间，并优先对这些点进行硬件测量。
• 利用聚类技术覆盖整个子空间，从而在保持探索质量的同时减少昂贵的真实硬件测量次数。
• 采用基于实测推理时间的奖励函数，引导 RL 智体选择运行更快的优化代码。
• 结合 RL 策略学习与成本模型，通过聚焦于搜索空间中信息量高的区域，避免冗余测量。
• 将该框架实现为 ReLeASE，一个与 TVM 及其他 DNN 框架兼容的开源优化编译器。

实验结果

研究问题

• RQ1 强化学习是否能在 DNN 编译优化中超越传统随机优化方法（如模拟退火和遗传算法）？
• RQ2 基于聚类的自适应采样是否能在不牺牲解质量的前提下减少所需硬件测量次数？
• RQ3 将强化学习与自适应采样结合，是否能实现更快收敛并生成更优的代码？
• RQ4 自适应采样在已使用成本模型的框架（如 AutoTVM）中，能在多大程度上带来改进？

主要发现

• 在 ResNet-18 上，与 AutoTVM 相比，ReLeASE 将优化时间减少 4.45 倍，分别为 1.20 小时与 9.13 小时。
• 在 AlexNet、VGG-16 和 ResNet-18 上，该方法将推理性能最高提升 5.6%，端到端推理时间从 1.0277 ms 降低至 0.9673 ms（AlexNet）。
• 自适应采样将 AutoTVM 的模拟退火优化时间提升 4.00 倍，证明其具备良好的通用性。
• 基于 RL 的方法在更少的步数内收敛，且以更少的搜索努力实现更优或相当的性能。
• 自适应采样算法通过聚焦于聚类生成的代表性子空间，减少了硬件测量次数，提升了采样效率。
• ReLeASE 作为开源软件公开发布，可无缝集成至现有 DNN 编译流水线。