如何确保 tensorflow 正在为一个会话使用所有 cpu 个核心？

Question

我想部署一个带有 tensorflow-serving 的 keras 模型。该模型从 keras .h5 模型转换为 .pb 文件。（原模型来自[这里][https://github.com/shaoanlu/face_toolbox_keras））

在此模型上使用 keras 进行推理时，如果我只使用我的 CPU，则 12 个内核都在工作，推理平均需要 0.7 秒。

转换模型并使用tensorflow serving时，只用了一个核，平均耗时2.7s。

我尝试将 --tensorflow_session_parallelism、--tensorflow_intra_op_parallelism 和 --tensorflow_inter_op_parallelism 等选项设置为 12，但没有任何变化，从 tfserving 容器内部查看顶部时只有一个核心在工作。

我还尝试为我的机器架构编译 tensorflow-serving，我得到了轻微的改进（2.7 秒到 2.5 秒），但我无法控制每个会话使用的内核数量。

我认为其他内核可用于并发请求是件好事，但我想拥有更多控制权。

Answer 1

问题可能是由于 constant folding pass。使用 tf.placeholder 应该可以解决问题。

                if args.const_fold:
                    A = tf.ones([size, size], name=("A%s" % i))
                    B = tf.ones([size, size], name=("B%s" % i))
                else:
                    A_name = "A%s" % i
                    B_name = "B%s" % i
                    A = tf.placeholder(tf.float32, shape=[size, size], name=A_name)
                    B = tf.placeholder(tf.float32, shape=[size, size], name=B_name)
                    feed["%s:0" % A_name] = np.random.rand(size, size)
                    feed["%s:0" % B_name] = np.random.rand(size, size)

据我了解，您的代码可能如上图if block所示。将其更改为 else block 应该可以解决问题。

更多信息，请参考此

如何确保 tensorflow 正在为一个会话使用所有 cpu 个核心？

How to make sure tensorflow is using all cpu cores for a session?

python

tensorflow

keras

tensorflow-serving