带有XLA的Tensorflow没有充分利用CPU容量

Question

我创建了一个在Tensorflow 2.5中实现的蒙特卡罗模拟模型。该模型主要由tf.while_loop内部的向量乘法组成。我在一台有8个虚拟cpu的Linux机器上对性能进行基准测试。当我在图形模式下运行模型(没有XLA优化)时，该模型充分利用了所有8个CPU(使用top命令我可以看到%CPU接近800%)。然而，当我用XLA编译后运行模型时(通过在@tf.function装饰器中使用jit_compile=True)，我可以看到%CPU利用率接近250%。有没有办法强制Tensorflow使用XLA来利用所有可用的CPU容量?

我已经尝试改变inter_op_parallelism和intra_op_parallelism设置。虽然将这两个线程设置为1可以将CPU利用率从250%降低到100%，但将它们设置为8不会将利用率提高到250%以上。

有什么帮助和建议，可能会发生什么?

Answer 1

我也有同样的问题。使用这里找到的建议:https://www.tensorflow.org/xla我将ML模型的JIT编译序列修改为类似

的内容

os.environ['XLA_FLAGS'] = '--xla_dump_to=/tmp/dump'
@tf.function(jit_compile=True)
def foo(data):
return model(data)

这在/tmp/dump中产生了一个object (*.o)文件，我用objdump -d对它进行了反汇编。看看反汇编，编译器似乎为模型和计算内核生成了直线代码，而不是调用可能支持并行执行的库。我没有看到任何迹象表明并行执行这个jit模型的可能性，尽管像你一样，当我简单地调用模型时，我确实观察到并行执行。

然而，对我来说，这个特定模型的最佳性能来自于使用@tf.function()和jit_compile=False。在这种情况下，我观察到' intr_op '并行发生-但没有'inter_op'并行发生，这在简单调用模型时也会观察到。