Tensorflow XLA会让它变慢吗?

时间:2017-06-13 14:02:08

标签: tensorflow gpu

我正在编写一个启用了XLA的非常简单的tensorflow程序。基本上它是这样的:

import tensorflow as tf

def ChainSoftMax(x, n)
  tensor = tf.nn.softmax(x)
  for i in range(n-1):
    tensor = tf.nn.softmax(tensor)
  return tensor

config = tf.ConfigProto()
config.graph_options.optimizer_options.global_jit_level = tf.OptimizerOptions.ON_1

input = tf.placeholder(tf.float32, [1000])
feed = np.random.rand(1000).astype('float32')

with tf.Session(config=config) as sess:
  res = sess.run(ChainSoftMax(input, 2000), feed_dict={input: feed})

基本上,我们的想法是查看XLA是否可以融合softmax链以避免多个内核启动。在启用XLA的情况下,上述程序比使用GPU卡的机器上的没有XLA的程序快2倍。在我的gpu配置文件中,我看到XLA生成了许多名为" reduce_xxx"的内核。和" fusion_xxx"这似乎压倒了整个运行时。谁知道这里发生了什么?

1 个答案:

答案 0 :(得分:0)

看看来自TF dev峰会的视频。他们使用微基准测试的图表显示了GPU中的XLA does not make everything faster

enter image description here

因此,如果某些操作的表现比没有XLA的情况差,那就不足为奇了。