簡單的TensorFlow示例將模型的一個副本加載到每個可用的GPU上

我正在尋找從不同的生產機器學習框架移植到TensorFlow。在我們當前的訓練和推理系統中，我們將模型的副本加載到機器上的GPU上。簡單的TensorFlow示例將模型的一個副本加載到每個可用的GPU上

我想現在保持這種負載平衡的方式。我在哪裏可以找到一個將TF模型的一個副本加載到機器上可用的每個GPU上的簡單示例？

2017-10-16 empty

下面是https://github.com/rafaljozefowicz/lm/blob/master/language_model.py#L21

一個例子你換你的模型創建代碼爲_forward函數，然後調用一次爲每個GPU

for i in range(hps.num_gpus): 
     with tf.device(assign_to_gpu(i, ps_device)), tf.variable_scope(tf.get_variable_scope(), 
                     reuse=True if i > 0 else None): 
      loss = self._forward(i, xs[i], ys[i], ws[i]) 
      losses += [loss] 
      if mode == "train": 
       cur_grads = self._backward(loss, summaries=(i == hps.num_gpus - 1)) 
       tower_grads += [cur_grads] 

    self.loss = tf.add_n(losses)/len(losses)

來源

2017-10-16 18:52:13

我想我不明白這一點的例子。它看起來像是在一個GPU上傳播，然後在另一個GPU上傳播。串行。我想要並行運行GPU上的模型。 – empty

這個例子向前傳播8個模型並行的8個gpus，然後加起來他們的損失 –

我在這裏沒有看到任何多線程。平行表現在哪裏？ – empty

簡單的TensorFlow示例將模型的一個副本加載到每個可用的GPU上

回答

相關問題