Tensorflow-如何在GPU之間拆分批次以在經過訓練的模型上進行預測？

Question

我正在使用未創建但已修改的模型（來自此倉庫https://github.com/GeorgeSeif/Semantic-Segmentation-Suite ）

我已經訓練過模型，可以使用它們進行足夠好的預測，但是我想遍歷整個圖像文件夾，並在多個GPU之間分配工作。 我不完全了解tf.device（）是如何工作的，而我嘗試過的工作根本沒有工作。

我以為我可以做這樣的事情：

for i, d in enumerate(['\gpu:0', '\gpu:1']):
    with tf.device(d):
        output = sess.run(network, feed_dict={net_input: image_batch[i]})

但這實際上並沒有將任務分配給不同的GPU，也沒有引發錯誤。

我的問題是，是否可以在不顯式修改網絡代碼預訓練的情況下將不同的圖像分配給單獨GPU上的會話的不同實例。 我想避免使用CUDA_VISIBLE_DEVICES = ...來運行兩個不同的python腳本。

有沒有簡單的方法可以做到這一點？

據我了解，操作的定義必須嵌套在“ with tf.device（）”塊中，但是，在推斷操作時，只是模型和權重的加載，但是如果我將其放入“ with tf。”中。 device（）“塊，我收到一條錯誤消息，說該圖已經存在，無法定義兩次。

Answer 1

tf.device僅在構建圖而不執行圖時適用，因此將session.run調用包裝在設備上下文中不會執行任何操作。

相反，我建議您使用tf復制器或tf分發策略（取決於tf版本的tf.distribute / tf.contrib.distribute），特別是MirroredStrategy。