[论文解读] REGAL: Transfer Learning For Fast Optimization of Computation Graphs.

Aditya Paliwal, Felix Gimeno+5 more|arXiv (Cornell University)|May 7, 2019

Parallel Computing and Optimization Techniques参考文献 20被引用 12

一句话总结

REGAL 提出了一种基于深度强化学习的迁移学习方法，用于优化神经网络计算图，以减少执行时间和峰值内存占用。它在离线阶段训练一个优化器，使其能泛化到未见过的计算图上，实现在数秒内获得高质量优化结果——远快于先前方法所需的数小时——同时优于经典方法和基于学习的基线方法。

ABSTRACT

We present a deep reinforcement learning approach to minimizing the execution cost of neural network computation graphs in an optimizing compiler. Unlike earlier learning-based works that require training the optimizer on the same graph to be optimized, we propose a learning approach that trains an optimizer offline and then generalizes to previously unseen graphs without further training. This allows our approach to produce high-quality execution decisions on real-world TensorFlow graphs in seconds instead of hours. We consider two optimization tasks for computation graphs: minimizing running time and peak memory usage. In comparison to an extensive set of baselines, our approach achieves significant improvements over classical and other learning-based methods on these two tasks.

研究动机与目标

• 解决现有优化方法在面对新图时需重新训练的问题。
• 实现在未见过模型上无需微调即可快速获得高质量的计算图优化。
• 最小化神经网络计算图中的执行时间和峰值内存使用量。
• 使单一训练好的优化器能泛化到多种真实世界的 TensorFlow 计算图上。

提出的方法

• 在多样化计算图上离线训练一个深度强化学习智能体，以学习优化策略。
• 智能体学习做出诸如算子融合、布局转换和内存分配等决策。
• 应用迁移学习，使同一策略网络能泛化到此前未见过的图上，而无需进一步训练。
• 该方法使用结合执行时间和内存使用量的奖励函数，以指导优化决策。
• 采用离策略强化学习与经验回放训练策略网络，以提高样本效率。
• 将该框架集成到优化编译器流水线中，以便在推理时应用决策。

实验结果

研究问题

• RQ1 是否可以仅训练一次强化学习策略，并将其泛化应用于优化此前未见过的计算图？
• RQ2 与经典优化技术相比，该方法在执行时间和内存使用量方面表现如何？
• RQ3 迁移学习在多大程度上可减少优化新模型时的重训练需求？
• RQ4 该方法是否能在显著缩短优化时间的同时，保持或提升优化质量？

主要发现

• REGAL 在执行时间和峰值内存使用量方面，优化质量与经典方法和基于学习的基线方法相当或更优。
• 该方法将真实世界 TensorFlow 图的优化时间从数小时缩短至数秒。
• 训练好的策略无需微调即可有效泛化到未见过的图上，展现出强大的迁移能力。
• 该方法在多种模型上均优于强基线方法，能更有效地最小化执行时间和内存占用。
• 基于奖励的强化学习框架成功平衡了速度与内存效率之间的权衡。