Многопользовательский подход Tensorflow Java

У меня есть сервер с несколькими GPU и вы хотите в полной мере использовать их во время вывода модели внутри java-приложения. По умолчанию tenorflow захватывает все доступные графические процессоры, но использует только первый.

Я могу представить три варианта решения этой проблемы:

• Ограничить видимость устройства на уровне процесса, а именно с помощью переменной среды CUDA_VISIBLE_DEVICES.

  Это потребует от меня запуска нескольких экземпляров приложения java и распределения трафика между ними. Не эта соблазнительная идея.

• Запустите несколько сеансов внутри одного приложения и попробуйте назначить одно устройство каждому из них через 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.

• Создайте модель в режиме multi-tower в python, используя спецификацию tf.device(). На стороне 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) {
           // ...
       }
   }
   

   Как и другие упомянутые подходы, это не освобождает меня от ручного распределения данных между устройствами. Но, по крайней мере, он работает стабильно и сравнительно легко реализуется. В целом, это выглядит как (почти) нормальная техника.

Есть ли элегантный способ сделать такую ​​базовую вещь с помощью java API-интерфейса tensorflow? Любые идеи были бы хорошы.