Tensorflow Java多GPU推理

Question

我有一個具有多個GPU的服務器，並希望在Java應用程序內的模型推理期間充分利用它們。 默認情況下，tensorflow會占用所有可用的GPU，但僅使用第一個GPU。

我可以想出三個選項來克服這個問題：

1. 在進程級別限制設備可見性，即使用CUDA_VISIBLE_DEVICES環境變量。

   這將要求我運行java應用程序的幾個實例並在它們之間分配流量。 不是那個誘人的想法。

2. 在單個應用程序中啟動多個會話，並嘗試通過ConfigProto為每個會話分配一個設備：

    public class DistributedPredictor { private Predictor[] nested; private int[] counters; // ... public DistributedPredictor(String modelPath, int numDevices, int numThreadsPerDevice) { nested = new Predictor[numDevices]; counters = new int[numDevices]; for (int i = 0; i < nested.length; i++) { nested[i] = new Predictor(modelPath, i, numDevices, numThreadsPerDevice); } } public Prediction predict(Data data) { int i = acquirePredictorIndex(); Prediction result = nested[i].predict(data); releasePredictorIndex(i); return result; } private synchronized int acquirePredictorIndex() { int i = argmin(counters); counters[i] += 1; return i; } private synchronized void releasePredictorIndex(int i) { counters[i] -= 1; } } public class Predictor { private Session session; public Predictor(String modelPath, int deviceIdx, int numDevices, int numThreadsPerDevice) { GPUOptions gpuOptions = GPUOptions.newBuilder() .setVisibleDeviceList("" + deviceIdx) .setAllowGrowth(true) .build(); ConfigProto config = ConfigProto.newBuilder() .setGpuOptions(gpuOptions) .setInterOpParallelismThreads(numDevices * numThreadsPerDevice) .build(); byte[] graphDef = Files.readAllBytes(Paths.get(modelPath)); Graph graph = new Graph(); graph.importGraphDef(graphDef); this.session = new Session(graph, config.toByteArray()); } public Prediction predict(Data data) { // ... } } 

   這種方法似乎一目了然。 但是，會話偶爾會忽略setVisibleDeviceList選項，並且所有會話都會導致第一個設備導致Out-Of-Memory崩潰。

3. 使用tf.device()規范在python中以多塔方式構建模型。 在java端，在共享會話中給出不同的Predictor不同的塔。

   對我來說，感到麻煩和慣用錯誤。

更新：正如@ash提出的那樣，還有另一種選擇：

4. 通過修改其定義（ graphDef ）為現有圖的每個操作分配適當的設備。

   要完成它，可以調整方法2中的代碼：

    public class Predictor { private Session session; public Predictor(String modelPath, int deviceIdx, int numDevices, int numThreadsPerDevice) { byte[] graphDef = Files.readAllBytes(Paths.get(modelPath)); graphDef = setGraphDefDevice(graphDef, deviceIdx) Graph graph = new Graph(); graph.importGraphDef(graphDef); ConfigProto config = ConfigProto.newBuilder() .setAllowSoftPlacement(true) .build(); this.session = new Session(graph, config.toByteArray()); } private static byte[] setGraphDefDevice(byte[] graphDef, int deviceIdx) throws InvalidProtocolBufferException { String deviceString = String.format("/gpu:%d", deviceIdx); GraphDef.Builder builder = GraphDef.parseFrom(graphDef).toBuilder(); for (int i = 0; i < builder.getNodeCount(); i++) { builder.getNodeBuilder(i).setDevice(deviceString); } return builder.build().toByteArray(); } public Prediction predict(Data data) { // ... } } 

   就像其他提到的方法一樣，這個方法並沒有讓我免於在設備之間手動分配數據。 但至少它運行穩定，並且相對容易實現。 總的來說，這看起來像（幾乎）正常的技術。

使用tensorflow java API有一種優雅的方式來做這樣的基本事情嗎？ 任何想法，將不勝感激。

Answer 1

簡而言之：有一種解決方法，每個GPU最終會有一個會話。

細節：

一般流程是TensorFlow運行時尊重為圖中的操作指定的設備。 如果沒有為操作指定設備，則它會根據某些啟發式“放置”它。 這些啟發式技術目前導致“在GPU上進行操作：0如果GPU可用並且有GPU內核用於操作”（ Placer::Run以防您感興趣）。

您認為我要求的是TensorFlow的合理功能請求 - 能夠將序列化圖形中的設備視為“虛擬”設備，在運行時映射到一組“phyiscal”設備，或者設置“默認設備” ”。 此功能目前不存在。 向ConfigProto添加此選項是您可能要為其提交功能請求的內容。

我可以在此期間提出一個解決方法。 首先，對您提出的解決方案進行一些評論。

1. 你的第一個想法肯定會奏效，但正如你所指出的那樣，很麻煩。

2. 在ConfigProto中使用visible_device_list進行設置並不ConfigProto ，因為這實際上是一個每進程設置，並且在進程中創建第一個會話后被忽略。 這肯定沒有記錄，應該是（並且有點不幸的是，它出現在每會話配置中）。 但是，這解釋了為什么您的建議不起作用以及您仍然看到使用單個GPU的原因。

3. 這可行。

另一個選擇是最終得到不同的圖形（操作明確放在不同的GPU上），每個GPU產生一個會話。 這樣的東西可以用來編輯圖形並明確地為每個操作分配一個設備：

public static byte[] modifyGraphDef(byte[] graphDef, String device) throws Exception {
  GraphDef.Builder builder = GraphDef.parseFrom(graphDef).toBuilder();
  for (int i = 0; i < builder.getNodeCount(); ++i) {
    builder.getNodeBuilder(i).setDevice(device);
  }
  return builder.build().toByteArray();
} 

之后，您可以使用以下內容為每個GPU創建一個Graph和Session ：

final int NUM_GPUS = 8;
// setAllowSoftPlacement: Just in case our device modifications were too aggressive
// (e.g., setting a GPU device on an operation that only has CPU kernels)
// setLogDevicePlacment: So we can see what happens.
byte[] config =
    ConfigProto.newBuilder()
        .setLogDevicePlacement(true)
        .setAllowSoftPlacement(true)
        .build()
        .toByteArray();
Graph graphs[] = new Graph[NUM_GPUS];
Session sessions[] = new Session[NUM_GPUS];
for (int i = 0; i < NUM_GPUS; ++i) {
  graphs[i] = new Graph();
  graphs[i].importGraphDef(modifyGraphDef(graphDef, String.format("/gpu:%d", i)));
  sessions[i] = new Session(graphs[i], config);    
}

然后使用sessions[i]在GPU #i上執行圖形。

希望有所幫助。