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端，在共享会话中给不同的Predictors不同的塔。 对我来说，感到麻烦和惯用错误。

更新：正如@a​​sh提出的那样，还有另一种选择：

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. 在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上执行图形。

希望有所帮助。