我创建了一个在 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% 以上。
关于可能发生的事情的任何帮助和建议?