学界 | 谷歌提出基于强化学习的优化配置方法：可让TensorFlow更充分利用CPU和GPU

来源:机器之心

原标题：学界 | 谷歌提出基于强化学习的优化配置方法：可让TensorFlow更充分利用CPU和GPU

选自arXiv

作者：Azalia Mirhoseini等

机器之心编译

参与：吴攀、李泽南

众所周知，深度学习是非常计算密集的，合理分配计算资源对于提升运算速度、节省能源等有着至关重要的价值。近日，谷歌的一篇 ICML 2017 论文提出了一种使用强化学习优化设备配置的方法，据称可以在一些网络计算上实现优于专家优化的方法的表现。机器之心对本文进行了摘要介绍。

论文：使用强化学习的设备配置优化（Device Placement Optimization with Reinforcement Learning）

论文地址：https://arxiv.org/abs/1706.04972

摘要

过去几年，我们见证了神经网络的训练和推理推动着计算能力需求的增长。目前解决这种需求的常用办法是搭建异构分布式计算环境，其中包含 CPU 和 GPU 这样的硬件设备。重要的是，神经网络模型在设备上的计算资源分配通常是由人类专家依靠简单的探索和直觉做出的。在本文中，我们提出了一种为 TensorFlow 计算图学习并优化设备任务分配的方法。其中最关键的思路是使用序列到序列模型来预测 TensorFlow 图中哪些子集应该运行在哪些可用设备上，然后将预测分配的执行时间用作奖励信号以优化序列到序列模型的参数。我们的结果在用于 ImageNet 分类的 Inception-V3 任务、RNN LSTM 语言建模和神经机器翻译任务上的表现优于普通手动设备配置（device placement）方式。

图 1 ：基于强化学习的设备配置模型概览

图 2：设备配置模型的架构

3.2 架构细节

我们使用了一种序列到序列模型（Sutskever et al., 2014）来预测设备配置，该模型带有 LSTM（Hochreiter & Schmidhuber, 1997）和一个基于内容的注意机制（Bahdanau et al., 2015）。图 2 展示了我们模型的整体架构，其可以分成两个部分：编码器 RNN 和解码器 RNN。

编码器 RNN 的输入是输入图（input graph）的指令的序列。我们通过连接它们的信息嵌入这些指令。具体来说，对于每一个图 G，我们首先收集其指令的类型。其指令的类型描述了其基本的计算，比如 MatMul 或 conv2d。对于每一种类型，我们都存储一个可调节的嵌入向量。然后我们记录输出张量的每个指令列表的大小，然后把它们连接成一个固定大小的填充了 0 的列表，atv，称为输出形状（output shape）。我们也取其 one-hot 编码向量，其表征了那些将输入和输出指引到每个指令的指令。最后，每个指令的嵌入都是其类型、其输出形状和其 one-hot 编码的邻接信息的连接。

该解码器是一个带有固定时间步骤数量的 attentional LSTM（Bahdanau et al., 2015），这个数量等于图 G 中的指令的数量。在每一步，该解码器输出该设备（device）以运行在同一编码器时间步骤的指令。每个设备都有其自己的可调节嵌入，然后其被用作下一个解码器时间步骤的输入。

图 3：分布式和异步的参数更新与奖励评估

表 2：由基于强化学习的方法所找到的设备配置方式和基准方法的运行时间比较（单位：秒，越低越好）。

对于每一个模型，第一行给出了带有 1 个 CPU 和 2 个 GPU 的结果；第二行给出了带有 1 个 CPU 和 4 个 GPU 的结果。最后一列给出了基于强化学习的配置优于最快基准方法的运行时间的提升。为了减少方差，少于 10 秒的运行时间都测试了 15 次，这里给出的是平均值。OOM 是指内存不足（Out Of Memory）。

图 4：神经机器翻译图的基于强化学习的配置。上面是编码器，开奖，下面是解码器。设备用不同颜色表示，其中无色表示指令在 CPU 上运行，不同的其它颜色表示不同的 GPU。相比于专家精心设计的配置，这个配置在运行时间上实现了 19.3% 的提升。