如何將TensorFlow圖形（模型）分割到多個GPU上以避免OOM？

Question

因此，我用TensorFlow r1.2實現了這個非常大而深的模型，運行在具有12 GB內存的NVIDIA Tesla k40上。該模型由多個RNN，一堆權重和嵌入矩陣以及偏向量組成。當我啓動培訓計劃時，首先需要大約2-3小時才能建模，然後由於OOM問題而崩潰。我試圖將批量減少到每批次的1個數據樣本，但仍遇到同樣的問題。

如果我谷歌tensorflow muitlple gpu，這些例子我發現主要集中在通過平行模型設計，這意味着有每個GPU運行相同圖形和具有CPU計算總梯度從而傳播利用多個GPU回到每個參數。

我知道一個可能的解決方案可能是在內存較大的GPU上運行該模型。但我不知道是否有辦法拆分我的圖形（模型）分成不同的部分按順序並將它們分配給不同的GPU？

Answer 1

關於using GPUs的官方指南向您展示了「使用多個GPU」中的示例。您只需在不同的tf.device上下文中創建操作;節點仍然會被添加到相同的圖表中，但是它們將被標註有設備指令，指示它們應該在哪裏運行。例如：

with tf.device("/gpu:0"): 
    net0 = make_subnet0() 
with tf.device("/gpu:1"): 
    net1 = make_subnet1() 
result = combine_subnets(net0, net1)