tensorflow - 带有 XLA 的 TensorFlow 没有充分利用 CPU 容量

Question

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

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

关于可能发生的事情的任何帮助和建议？

标签： tensorflowtensorflow2.0xlatensorflow-xla