如何在tensorflow 1.14.0中有效地使用OpenMP并行性

Question

我目前正在尝试寻找一种有效的方法，使用tensorflow在一定数量的内核上运行机器学习任务。从这些信息中，我发现有两种主要的方法可以做到这一点。

第一个是使用两个tensorflow变量intra_op_parallelism_threads和inter_op_parallelism_threads，然后使用此配置创建一个会话。

第二个是使用OpenMP。通过设置环境变量OMP_NUM_THREADS，可以操作为进程派生的线程数量。

当我发现通过conda和pip安装tensorflow会产生两种不同的环境时，我的问题就出现了。在conda install中，修改OpenMP环境变量似乎改变了进程并行化的方式，而在'pip环境‘中，唯一改变它的似乎是我前面提到的内部/内部配置变量。

由于基准测试的原因，这导致在尝试比较这两个安装时遇到了一些困难。如果我在conda install上的48核处理器上将OMP_NUM_THREADS设置为1，并将inter/intra设置为16，那么我只能得到大约200%的CPU使用率，因为大多数线程在任何给定时间都是空闲的。

omp_threads = 1
mkl_threads = 1
os.environ["OMP_NUM_THREADS"] = str(omp_threads)
os.environ["MKL_NUM_THREADS"] = str(mkl_threads)

config = tf.ConfigProto()
config.intra_op_parallelism_threads = 16
config.inter_op_parallelism_threads = 16
session = tf.Session(config=config)
K.set_session(session)

我预计这段代码会产生32个线程，其中大多数线程在任何给定的时间都会被利用，而实际上它会产生32个线程，而一次只有4-5个线程在使用。

以前有人在使用tensorflow时遇到过类似的情况吗？

为什么通过conda安装和通过pip安装会产生两种不同的环境？

通过使用前面讨论的两种方法的某种组合，有没有办法在两个安装上具有相当的性能？

最后，有没有更好的方法将python限制在特定数量的内核中？

Answer 1

回答你的第一个也是最后一个问题。

是的，我在使用通过pip安装的TensorFlow时遇到了类似的情况。在linux上，您可以使用线程亲和性、numatcl或taskset将python限制为特定数量的核心。

查看以下链接提供的详细信息，TensorFlow将始终生成多个线程，并且大多数线程将默认处于休眠状态。

• How to stop TensorFlow from multi-threading
• How can I run Tensorflow on one single core?
• Change number of threads for Tensorflow inference with C API

Answer 2

我认为这里的重点是conda使用MKL安装Tensorflow，而pip不这样做。

OpenMP控制仅适用于MKL，并且在pip安装中，OpenMP环境变量不起作用，仅将TFSessionConfig设置为intra/inter可能会影响多线程。