最近在实现A3C论文【1】算法的过程中,发现目前目前网上还没有太多资料讲解如何进行梯度累积,对于tensorflow分布式计算的异步更新也没有实验论证。因此将自己做的一点研究整理出来,还请各路大神指正。
一、问题描述
在Asynchronous methods中,使用了target network以避免网络变化太快。每个learner在一个训练epoch开始时会拷贝target network的权值, 训练一段时间后将梯度累积并用之更新target network,之后结束这个epoch。以下是n-step Q-learning的更新方式:
因此,在实现中,首先需要让各个learner能够获得target network的权值,然后要实现learner内部的梯度累积,最后要将这个梯度返回target network。
二、梯度计算、累积与更新
将梯度在target network和learner间传递的功能在distributed tensorflow中默认已经实现好了。Between-graph的方式中,每个thread会拷贝一份Graph,计算后回传回主Graph。需要解决的主要是梯度累积的问题。
基本的思路是:
repeat:
计算梯度
存储梯度
until 一定次数
将累积的梯度回传至target network
1. 定义操作
# Define input and output
with tf.name_scope('input'):
x = tf.placeholder(tf.float32, name="x")
with tf.name_scope('weights'):
w = tf.Variable(2.0, name='target_w')
with tf.name_scope('output'):
y = tf.mul(x, w, name='y')
with tf.name_scope('real_output'):
y_ = tf.placeholder(tf.float32, name="y_")
# Define train op
with tf.name_scope('train'):
optimizer = tf.train.GradientDescentOptimizer(LEARNING_RATE)
with tf.name_scope('gradient'):
loss = tf.reduce_mean(tf.square(y_ - y)) # MSE loss
gradient_all = optimizer.compute_gradients(loss) # gradient of network (with NoneType)
grads_vars = [v for (g, v) in gradient_all if g is not None] # all variable that has gradients
gradient = optimizer.compute_gradients(loss, grads_vars) # gradient of network (without NoneType)
grads_holder = [(tf.placeholder(tf.float32, shape=g.get_shape()), v)
for (g, v) in gradient]
train_op = optimizer.apply_gradients(grads_holder)
y_是真实值,y是网络的输出,loss是mse损失。optimizer是一个梯度下降优化器。gradient_all是optmizer计算的梯度,返回的是一个列表,其中的元素是型为的tuple。注意,如果网络不是单输入单输出(例如ac网络中有两个输出),那么compute_gradients可能会返回(None,v),即部分变量没有对应的梯度,在下一步的时候NoneType会导致错误。因此,需要将有梯度的变量提取出来,记为grads_vars。之后,对grads_vars再一次计算梯度,得到了gradient。最后, 生成一个placeholder用于存储梯度,并调用apply_gradients更新网络。
注意,此处定义的会出现问题,在第三节实验部分会提出一个解决办法。
2. 应用操作
# calculate gradients every
grads = []
for i in range(THREAD_STEPS):
x_i = ...
y_real = ...
y_i = sess.run(y, feed_dict={x: x_i})
loss_i = sess.run(loss, feed_dict={x: x_i, y_: y_real})
grad_i = sess.run(gradient, feed_dict={x: x_i, y_: y_real})
grads.append(grad_i)
# calculate total gradients
grads_sum = {}
# add up dθ
for i in range(len(grads_holder)):
k = grads_holder[i][0]
grads_sum[k] = sum([g[i][0] for g in grads])
# Apply gradients
_ = sess.run(train_op, feed_dict=grads_sum)
操作分为三步:
第一步,输入x_i,y_计算梯度,并使用一个列表grads保存梯度;
第二步,使用一个字典对梯度进行累积。字典的格式为。由于grads的每一个元素都是与grads_holder格式相同的梯度列表,grads_holder[i][0]对应的梯度值列表就是[g[i][0] for g in grads]。
第三步,将前一步生成的字典feed给apply_gradients,Mission complete。
三、实验设置与结果
简单验证了下以上方法的正确性。使用两个worker进行between-graph的异步训练,输入为[0,1,2],网络输出为,输出真实值定为10,损失为MSE。的初始值为2。优化器为梯度下降法,学习率为1。梯度计算公式为
我们设置了两个worker,各进行2 epoch*3 steps的训练。输出如下:
task0 - epoch0: x_i: 0. y_i: 0.0. loss: 100.0. grad: [(-0.0, 2.0)] task0 - epoch0: x_i: 1. y_i: 2.0. loss: 81.0. grad: [(-18.0, 2.0)] task0 - epoch0: x_i: 2. y_i: 4.0. loss: 64.0. grad: [(-32.0, 2.0)] task1 - epoch0: x_i: 0. y_i: 0.0. loss: 100.0. grad: [(-0.0, 2.0)] Final States of w in task0 - epoch0: 52.0
task0 - epoch1: x_i: 0. y_i: 0.0. loss: 100.0. grad: [(-0.0, 52.0)] task1 - epoch0: x_i: 1. y_i: 52.0. loss: 1681.0. grad: [(82.0, 52.0)]
...
首先,worker:0输入x:0,网络返回y:0,梯度计算为0,第二步输入x:1,y:2,梯度为-18,依次类推。在worker:0进行到1st epoch的第三步时,worker:1启动,注意此时w仍未2,网络没有发生改变,worker输入x:0后网络返回y:0 。
随后worker:0更新了梯度,梯度累积为, 新的权值更新为 。梯度累积的功能成功实现。
但是,我们对于graph间变量传递的机制理解有误,此时worker:1仍在第一个epoch,权值却也被更新为52(实际应保持该epoch开始时的2),本应为-18的梯度变成了82。
解决的办法是将thread和target network的权值分开。重新定义权值为:
with tf.name_scope('weights'):
target_w = tf.Variable(2.0, name='target_w')
w_list = [tf.Variable(2.0, name='target_w') for i in range(WORKER_NUM)]
w = w_list[FLAGS.task_index]
这样,我们创建了一个长度为thread数量的列表,用于存储各个进程的权值,w_list[task_index]是每个进程实际使用的权值,target_w是target network的权值。之后,定义权值更新的动作:
epoch_init = w.assign(target_w)
w_addup = tf.placeholder(tf.float32)
epoch_update = target_w.assign_add(w_addup)
在每次epoch开始前我们使用tf.assign(ref,value)将target_w的权值赋予w,epoch结束后将训练后权值与初始权值的差值增加给target_w。实验结果如下:
task0 - epoch0: x_i: 0. y_i: 0.0. loss: 100.0. grad: [(-0.0, 2.0)]
task0 - epoch0: x_i: 1. y_i: 2.0. loss: 81.0. grad: [(-18.0, 2.0)]
task1 - epoch0: x_i: 0. y_i: 0.0. loss: 100.0. grad: [(-0.0, 2.0)]
task0 - epoch0: x_i: 2. y_i: 4.0. loss: 64.0. grad: [(-32.0, 2.0)]
task1 - epoch0: x_i: 1. y_i: 2.0. loss: 81.0. grad: [(-18.0, 2.0)]
Final States of w in task0 - epoch0: 52.0
task0 - epoch1: x_i: 0. y_i: 0.0. loss: 100.0. grad: [(-0.0, 52.0)]
task1 - epoch0: x_i: 2. y_i: 4.0. loss: 64.0. grad: [(-32.0, 2.0)]
...
可以看到,worker:0已经完成了一次更新,梯度累积为-50,在epoch1中,初试的权值已经变为52。而worker:1的权值还是保持为2.
...
task0 - epoch1: x_i: 1. y_i: 52.0. loss: 1681.0. grad: [(82.0, 52.0)]
Final States of w in task1 - epoch0: 52.0
task0 - epoch1: x_i: 2. y_i: 104.0. loss: 8464.0. grad: [(368.0, 52.0)]
Final States of w in task0 - epoch1: -398.0
Final target_w: -348.0
done
task1 - epoch1: x_i: 0. y_i: 0.0. loss: 100.0. grad: [(-0.0, 102.0)]
task1 - epoch1: x_i: 1. y_i: 102.0. loss: 8281.0. grad: [(182.0, 102.0)]
task1 - epoch1: x_i: 2. y_i: 204.0. loss: 36864.0. grad: [(768.0, 102.0)]
Final States of w in task1 - epoch1: -848.0
Final target_w: -1298.0
done
至此,我们完成了梯度累积的异步更新的全过程。完整的Gist:allenwoods/async_grad_verify.py
引用:
[1] V. Mnih et al., “Asynchronous methods for deep reinforcement learning,” arXiv preprint arXiv:1602.01783, 2016.
[2] Tensorflow Python API : processing-gradients-before-applying-them
Tips:
1. 多进程使用GPU会导致OUT_OF_MEMORY_ERROR,这是由于tf默认会给任一进程分配所有能分配的显存,这样除了第一个进程其他进程都无显存可用。解决办法有两个,一是在运行命令前添加 CUDA_VISIBLE_DEVICES=9999(或任一大于你的显卡数的数字)禁用显卡,推荐对ps进程使用。二是在server配置里添加gpu_options=tf.GPUOptions(allow_growth=True)(或gpu_fraction)使得tf不会一次将显存分配完而是随着使用量逐渐增加,具体在以上提供的gist中有例子。
2. 运行命令为
python async_grad_test.py --ps_hosts=0.0.0.0:53198 --worker_hosts=0.0.0.0:58557,0.0.0.0:42832 --job_name=ps --task_index=0
python async_grad_test.py --ps_hosts=0.0.0.0:53198 --worker_hosts=0.0.0.0:58557,0.0.0.0:42832 --job_name=worker --task_index=0
python async_grad_test.py --ps_hosts=0.0.0.0:53198 --worker_hosts=0.0.0.0:58557,0.0.0.0:42832 --job_name=worker --task_index=1
3. ps进程通过Ctrl+c的方式无法关闭,需要通过kill的方式杀掉。批量关闭所有训练进程可使用以下命令:
ps -ef | grep /opt/anaconda3/bin/python| grep async_grad | awk {'print $2'} | xargs kill
/opt/anaconda3/bin/python 是Python解释器,async_grad是运行py文件的关键字,根据你的具体情况进行修改。
4. 运行Gist前需要先删除checkpoint中之前的记录,否则tf会认为任务已经完成从而不进行任何操作。