中文

OLLIE:基于推导的张量程序优化器

机器学习 2022-08-04 v1 性能

摘要

由于深度神经网络(DNN)在真实任务中的广泛采用,提升其运行时性能至关重要。现有的优化 DNN 张量代数表达式的方法仅考虑可由一组预定义算子表示的固定表达式,遗漏了通用表达式之间可能的优化机会。我们提出 OLLIE,首个基于推导的张量程序优化器。OLLIE 通过利用通用张量代数表达式之间的变换来优化张量程序,实现了显著更大的表达式搜索空间,其中包含先前工作支持的表达式作为特例。OLLIE 使用混合的基于推导的优化器,有效结合探索性与引导式推导,以快速发现高度优化的表达式。在七个 DNN 上的评估表明,OLLIE 在 A100 GPU 上可超越现有优化器至多 2.73×\times(平均 1.46×\times),在 V100 GPU 上至多 2.68×\times(平均 1.51×\times)。

关键词

引用

@article{arxiv.2208.02025,
  title  = {OLLIE: Derivation-based Tensor Program Optimizer},
  author = {Liyan Zheng and Haojie Wang and Jidong Zhai and Muyan Hu and Zixuan Ma and Tuowei Wang and Shizhi Tang and Lei Xie and Kezhao Huang and Zhihao Jia},
  journal= {arXiv preprint arXiv:2208.02025},
  year   = {2022}
}