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. 利用设置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();
} 

之后,您可以使用以下内容创建一个Graph和Session每个GPU:

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

希望有所帮助.