基于端到端可迁移深度强化学习的图优化

文 / Yanqi Zhou 和 Sudip Roy，研究员，Google Research

各式各样加速器上训练的大规模复杂神经网络驱动着越来越多的应用的发展。机器学习 (ML) 编译器为这一进程提供了便利，它将高级计算图映射到特定设备可执行的低级文件。为此，机器学习编译器需要解决许多优化问题，包括计算图重构、将算子分配给（不同的）计算设备、算子融合、张量的布局和平铺以及调度。例如，在训练设备配置优化问题 (Device Placement problem) 中，编译器需要确定计算图中的运算到目标物理设备之间的映射，以使目标函数（如训练步时间）最小化。

• 计算图
  https://www.easy-tensorflow.com/tf-tutorials/basics/graph-and-session#GRAPH

• 设备配置优化问题
  http://proceedings.mlr.press/v70/mirhoseini17a/mirhoseini17a.pdf

配置性能由多种复杂因素共同决定，包括设备间网络带宽、峰值设备内存、共置约束等，这对启发式算法或基于搜索的算法来说具有挑战性，而启发式算法或基于搜索的算法通常被用作快速但次优的最佳解决方案。此外，启发式算法很难开发和维护，尤其在新模型架构不断涌现的当下。

最近基于学习方法的研究成果展现出很好的前景，但也有一定局限性，难以在实践中部署。首先，这些方法不易泛化到未见过的计算图，特别是那些由较新的模型架构产生的计算图；其次，它们的样本效率很差，导致训练期间的资源消耗很高。最后，它们只能解决单一的优化任务，无法捕获编译栈中紧密耦合的优化问题的依赖关系。

在最近发表于 NeurIPS 2020 会议的口头报告论文“Transferable Graph Optimizers for ML Compilers”中，我们提出了一种用于计算图优化 (Graph Optimization, GO) 的端到端可迁移深度强化学习方法，该方法克服了上述所有限制。对照 TensorFlow 的默认优化，我们在三个计算图优化任务上展现出 33%-60% 的速度提升。在由多达 80,000 个节点（包括 Inception-v3、Transformer-XL 和 WaveNet）组成的多样化典型计算图集上，GO 比专家优化平均提高了 21%，比先前的最先进技术提高了 18%，收敛速度快了 15 倍。

• Transferable Graph Optimizers for ML Compilers

  https://papers.nips.cc/paper/2020/file/9f29450d2eb58feb555078bdefe28aa5-Paper.pdf

• TensorFlow
  https://tensorflow.google.cn/
  

  </