在多个GPU上运行Tensorflow模型推理脚本

Question

我正在尝试从tensorflow objec检测API运行模型评分（推理图）以在多个GPU上运行，尝试在主GPU中指定GPU编号，但仅在单个GPU上运行。placed GPU utilization snapshot here

使用tensorflow-gpu == 1.13.1，您能指出我在这里缺少的内容吗？

for i in range(2):
        with tf.device('/gpu:{}' . format(i)):
            tf_init()
            init = tf.global_variables_initializer
           with detection_graph.as_default():
              with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as session:
                   call to #run_inference_multiple_images function

Answer 1

对this question的回复应该为您提供一些解决方案。

除非另有说明，否则TensorFlow通常会占据所有可见的GPU。因此，如果您还没有尝试过，则可以删除with tf.device行（假设您只有两个GPU），而TensorFlow应该同时使用它们。

否则，我认为最简单的方法是使用os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"设置环境变量。