java - Tensorflow Java 多 GPU 推理

Question

我有一台带有多个 GPU 的服务器，我想在 Java 应用程序内的模型推理期间充分利用它们。默认情况下，tensorflow 占用所有可用的 GPU，但仅使用第一个。

我可以想到三个选项来解决这个问题:

1. 在进程级别限制设备可见性，即使用 CUDA_VISIBLE_DEVICES 环境变量。

   这将需要我运行 java 应用程序的多个实例并在它们之间分配流量。不是那种诱人的想法。

2. 在单个应用程序中启动多个 session ，并尝试通过 ConfigProto 为每个 session 分配一个设备:

   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) {
           // ...
       }
   }
   

   乍一看，这种方法似乎工作正常。但是， session 偶尔会忽略 setVisibleDeviceList 选项，并且所有这些都会针对第一个导致内存不足崩溃的设备。

3. 使用 tf.device() 规范在 python 中以多塔方式构建模型。在 Java 端，在共享 session 中为不同的 Predictor 提供不同的塔。

   对我来说感觉很麻烦而且惯用错误。

更新:正如@ash 所建议的，还有另一种选择:

1. 通过修改其定义 (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 有一个 session 。

详细信息:

一般流程是 TensorFlow 运行时尊重为图中的操作指定的设备。如果没有为操作指定设备，则它会根据一些试探法“放置”它。这些启发式方法目前导致“在 GPU 上进行操作:0 如果 GPU 可用并且有用于操作的 GPU 内核”(Placer::Run 如果您感兴趣的话)。

我认为你要求的是对 TensorFlow 的合理功能请求 - 能够将序列化图中的设备视为“虚拟”设备以在运行时映射到一组“物理”设备，或者设置“默认设备”。此功能当前不存在。将这样的选项添加到 ConfigProto 是您可能想要提交功能请求的内容。

我可以在此期间建议一个解决方法。首先，对您提出的解决方案进行一些评论。

1. 您的第一个想法肯定会奏效，但正如您指出的那样，它很麻烦。

2. 在 ConfigProto 中使用 visible_device_list 进行设置并不完全可行，因为这实际上是每个进程的设置，在创建第一个 session 后会被忽略进行中。这当然没有按应有的方式记录(并且有点不幸，这出现在每 session 配置中)。但是，这解释了为什么您在此处的建议不起作用以及为什么您仍然看到正在使用单个 GPU。

3. 这可行。

另一种选择是最终得到不同的图形(操作明确地放置在不同的 GPU 上)，从而导致每个 GPU 一个 session 。像这样的东西可以用来编辑图形并为每个操作明确分配一个设备:

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

希望对您有所帮助。