Tensorflow中使用timeline来优化性能

  |   0 评论   |   0 浏览

背景

初体验

使用timeline来优化优化性能，timeline可以分析整个模型在forward和backward的时候,每个操作消耗的时间，由此可以针对性的优化耗时的操作。我之前尝试使用tensorflow多卡来加速训练的时候， 最后发现多卡速度还不如单卡快，改用tf.data来 加速读图片还是很慢，最后使用timeline分析出了速度慢的原因，timeline的使用如下

run_metadata = tf.RunMetadata()
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
config = tf.ConfigProto(graph_options=tf.GraphOptions(
        optimizer_options=tf.OptimizerOptions(opt_level=tf.OptimizerOptions.L0)))
with tf.Session(config=config) as sess:
    c_np = sess.run(c,options=run_options,run_metadata=run_metadata)
    tl = timeline.Timeline(run_metadata.step_stats)
    ctf = tl.generate_chrome_trace_format()
with open('timeline.json','w') as wd:
    wd.write(ctf)

然后到谷歌浏览器中打卡chrome://tracing 并导入 timeline.json ，最后可以看得如下图所示的每个操作消耗的时间，。

这里横坐标为时间，从左到右依次为模型一次完整的forward and backward过程中，每个操作分别在cpu,gpu 0, gpu 1上消耗的时间，这些操作可以放大，非常方便观察具体每个操作在哪一个设备上消耗多少时间。这里我们cpu上主要有QueueDequeue操作，这是进行图片预期过程，这个时候gpu在并行计算的所以gpu没有空等；另外我的模型还有一个PyFunc在cpu上运行，如红框所示，此时gpu在等这个结果，没有任何操作运行，这个操作应该要优化的。另外就是如黑框所示，gpu上执行的时候有很大空隙，如黑框所示，这个导致gpu上的性能没有很好的利用起来，最后分析发现是我bn在多卡环境下没有使用正确，bn有一个参数updates_collections我设置为None 这时bn的参数mean,var是立即更新的，也是计算完当前layer的mean,var就更新，然后进行下一个layer的操作，这在单卡下没有问题的， 但是多卡情况下就会写等读的冲突，因为可能存在gpu0更新（写）mean但此时gpu1还没有计算到该层，所以gpu0就要等gpu1读完mean才能写，这样导致了 如黑框所示的空隙，这时只需将参数设置成updates_collections=tf.GraphKeys.UPDATE_OPS 即可，表示所以的bn参数更新由用户来显示指定更新，如

update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
  with tf.control_dependencies(update_ops):
    train_op = optimizer.minimize(loss)

这样可以在每个卡forward完后，再更新bn参数，此时的写写不存在冲突。优化后，我的2卡训练获得了接近2倍的加速比。

参考