如何评价Google即将发表在ICML17上的Placement Optimization的论文?

论文链接如下:[1706.04972] Device Placement Optimization with Reinforcement Learning 接要: 过去几年,我们见证了神经网络的训练和推理推动着计算能力需求的增长。目前解决这种需求的常用办法是搭建异构分布式计算环境,其中包含 CPU 和 GPU 这样的硬件设备。重要的是,神经网络模型在设备上的计算资源分配通常是由人类专家依靠简单的探索和直觉做出的。在本文中,我们提出了一种为 TensorFlow 计算图学习并优化设备任务分…
关注者
514
被浏览
11361

这篇论文我看了差不多一天,一些粗浅的理解在这里和大家分享一下,也作为一种记录方式记录下来。水平有限,多多指教~

使用了增强学习的方法,以及sequence to sequence的模型也就是LSTM。整体框架如下:


placement作为RL框架的输入,Environment就是TensorFlow按照这种placement在实际系统中运行一下,其实际的运行时间runtime作为RL框架的输出。该输出和某个baseline进行计算得到Loss,根据得到的Loss优化改进placement再次输入给RL系统。

其中,Loss是由该次的runtime和历史所有平均runtime得到的,并且该历史平均runtime在计算时忽略掉第一次的runtime,因为第一次runtime的值一般比较不靠谱,可以当做是outlier。

而placement是由一个sequence to sequence模型得到的,也即是一个LSTM模型。该模型的框架如下:


模型输入:op的一些属性。图示中表明,该计算图一共100个op, 因此会有100个op相关的属性作为特征输入。op的属性包括以下三类:

(1)一个embedding向量,根据op类型比如MatMul(矩阵乘),conv2d(卷积)等,对这种特征进行embedding得到一个向量,可以通过参数调整改变embedding之后向量长度,因此是一个可变长的(tunable)向量。论文中并没有说明如何做embedding.

(2)一个定长的由output的shape拼接成的向量。拼接shape的时候按照output中shape维度最高的分配空间,例如,有两个output,其中一个shape是4维[20,40,20,30],另一个shape, 其中间2维[20, 30], 拼接在一起之后就是[20, 40, 20, 30, 0,20, 30,0],也就是论文中所说的fixed-size zero-padded.

(3) 一个用来表示op的输入输入的one hot encoding 向量。这也很明显是个定长的向量,统计计算图中一共有多少个op,并分别编号。然后对某op的输入输出进行one hot encoding编码,这个比较简单,不用解释。

模型输出:计算图中有多少个op,就有多少个输出。图示中就会有100个输出。

RL系统得到Loss之后如何优化现有模型,公式如下:



实际上,论文中使用了均值的方法代替公式中的期望值E:



公式的意思就是:根据误差改变参数theta, 再根据新的theta修改LSTM模型得到新的placement。其中R(Pi)表示某个placement的runtime,B表示历史平均Runtime。对于一些特殊情况,例如某个策略使得大部分op放置在同一个设备上导致内存不足,那么该次的runtime就设置为一个很大的常数,该常数一定是带来巨大的误差的。

以上就是整个论文中实验的框架,如果要进行该实验,需要准备n台机器,每台机器上有些设备,这些设备就是计算图op要进行放置的候选。但是这么做可能比较慢,因此论文中设置了分布式训练该RL模型的方案,即分布式训练。其架构如下:


图中每个的worker有多个device,这些device负责执行controller发送来的placement policy。训练过程就是:首先,每个Controller根据上面截图的公式(3)生成K个placement,这里的K是该controller控制的worker的个数,然后将这K个placement分别发给worker,每个worker一个placement, 每个worker就根据收到的placement设置op应该放在哪个设备上执行。其次,worker执行完之后,Controller收集对应的每个worker的runtime,每个worker的runtime也就是公式(3)的R(Pi)。这里有个小trick,在计算历史平均runtime时,第一个step的时间不计,并且是每10个step的runtime才计算一次R(Pi)。也就是RL系统每10个step才更新一次模型参数theta。

每个Controller的参数theta都存放在parameter server中,每次Controller根据公式(3)得到新梯度之后,都会和PS通信更新其参数,然后Controller再根据最新得到的参数做K次蒙特卡洛采样得到新的K个placement policy发给K个worker。

实验的结果,个人感觉并不明显,只有20%左右的提升。

以上就是论文的细节。

另外还有一点我觉得对实验影响较大,就是3.3节op的Co-locating那部分。由于有些模型的op个数比较多,并且由模型输入部分的组成可以看出,真实的输入是op个数乘以一个比较大的常数。这可能会导致模型训练时梯度vanishing难以收敛或者过拟合,因此论文提出由用户根据先验将某些op划分到一个group中,同一个group中的op放置在同一个设备上去。论文中给出了一个Co-locating前后的对比:



我之所以觉得这是个缺点,因为如果一个小白用户做了一些不合理的cl-locating,这将直接导致RL训练得不到好结果。主要是因为,用户一旦设置了group,该group内的op不能被拆解,只能merger新的op进入该op。

最后是关于效果的好坏,我觉得这篇论文一个比较大的意义就是,提供了一个用机器学习方法解决placement问题的思路,效果虽然没有想象的好但是也很有价值。还有就是关于并行策略的问题,我认为也是可以用这个方法学到对某个op做模型并行还是数据并行的,只需要在输入时加入并行策略的特征,输出时输出可能的并行策略就可以。