python - Tensorflow 2.0 未使用所有 CPU 内核

Question

我正在我的本地机器上使用 tensorflow 2.0 进行一些稀疏矩阵计算，该机器有 8 个内核。在执行操作时，我可以在 top 命令中看到只有 8 个内核中的 1 个被使用。

运算是稀疏乘法

tf.sparse.sparse_dense_matmul(a_tf_sp, b_tf_dense_T)

我如何告诉 tensorflow 2.0 使用所有核心进行此计算。

标签： pythontensorflowtensorflow2.0

Answer 1

你能在里面定义你的代码with tf.device('/cpu:0'):并检查吗？