在多 GPU 情况下，CPU 代码是否可以存在于“with tf.device(gpu_id):”下？

Question

您好，我是 Tensorflow 的新手，我被分配了一项任务来更改 Github 项目“tf-faster-rcnn”中的“Demo.py”以实现多 GPU 推理。

而这通常是我打算做的（假设我拥有的图像数量与 GPU 的数量相同，为了简单起见，我将使用此处未说明的队列）：

for id, gpu in gpu_dict:
    with tf.device(gpu):
        im_detect(images[id])

源文件中提供了“im_detect”函数（我可以直接调用它），它包含一些非GPU代码（如条件和数据准备）

def im_detect(sess, net, im):
  blobs, im_scales = _get_blobs(im)
  assert len(im_scales) == 1, "Only single-image batch implemented"

  im_blob = blobs['data']
  blobs['im_info'] = np.array([im_blob.shape[1], im_blob.shape[2], im_scales[0]], dtype=np.float32)

  _, scores, bbox_pred, rois = net.test_image(sess, blobs['data'], blobs['im_info'])

  boxes = rois[:, 1:5] / im_scales[0]
  scores = np.reshape(scores, [scores.shape[0], -1])
  bbox_pred = np.reshape(bbox_pred, [bbox_pred.shape[0], -1])
  if cfg.TEST.BBOX_REG:
    # Apply bounding-box regression deltas
    box_deltas = bbox_pred
    pred_boxes = bbox_transform_inv(boxes, box_deltas)
    pred_boxes = _clip_boxes(pred_boxes, im.shape)
  else:
    # Simply repeat the boxes, once for each class
    pred_boxes = np.tile(boxes, (1, scores.shape[1]))

  return scores, pred_boxes

由于我以前从未玩过GPU，而且我是Tensorflow的新手，所以我想问一下在Tensorflow中为每个GPU分配这样的函数调用是否可以？

----------------下面更新了------------------------------------

我知道 Tensorflow 中有一个“alow_soft_placement”选项，它将那些非 GPU 代码分配给 CPU，但是当有多个 GPU 时，一个 CPU 如何处理来自所有 GPU 的这些请求？我应该为每个 GPU 创建一个 CPU 线程吗？

Answer 1

是的。来自https://www.tensorflow.org/programmers_guide/using_gpu。如果操作没有 CUDA 内核，会话配置的 allow_soft_placement 参数允许 TensorFlow 回退到 CPU。

myConf = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)
sess = tf.Session(config=myConf)

有时你不会想要这个 - 例如。如果您正在尝试验证您期望的所有操作实际上都在 GPU 上运行。

您还可以在 with tf.device('/gpu:0'): 块内使用 with tf.device('/cpu:0'): 将操作显式分配给 CPU。

我倾向于使用严格放置，然后在 TensorFlow 抱怨时将不兼容的操作显式分配给 cpu。这样我就可以确定所有合适的操作都经过 GPU 优化。

更新：

这里有一些示意图代码，它应该概述如何在 GPU 上运行并行计算。

graph = tf.Graph()

with graph.as_default():

gpus = ['/gpu:0', '/gpu:1']
results = []
datasets = []

for idx, gpu in enumerate(gpus):
   with tf.device(gpu):
       # assign data prep ops to CPU
       # (or use soft placement and leave out the next line).
       with tf.device('/cpu:0'):
            datasets[idx] = tf.placeholder(tf.float32, name = 'Features'+idx)

       # Computationally expensive ops get assigned to GPU, but make reference
       # to specific non-GPU ops on CPU.
       results[idx] = tf.reduce_sum(datasets[idx])  

myConf = tf.ConfigProto(allow_soft_placement=True, log_device_placement=True)

with tf.Session(graph=graph, config=myConf) as session:

    # Now, using the graph set up previously, evaluate results
    # using both gpu devices (each these ops depends on independent
    # cpu ops).
    res0, res1 = session.run([results[0], results[1]])