Многопользовательский подход 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 есть еще один вариант:

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

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

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

11
13 дек. '17 в 21:35
источник поделиться
1 ответ

Вкратце: существует обходное решение, в результате которого вы получаете один сеанс на графический процессор.

Подробнее:

Общий поток - это то, что среда исполнения TensorFlow относится к устройствам, указанным для операций на графике. Если для операции не указано какое-либо устройство, оно "помещает" его на основании некоторых эвристик. Эти эвристики в настоящее время приводят к "операции места на графическом процессоре: 0, если доступны графические процессоры и есть ядро ​​GPU для операции" (Placer::Run в если вам интересно).

То, о чем вы просите, я думаю, это разумный запрос функции для TensorFlow - возможность рассматривать устройства в сериализованном графике как "виртуальные", которые должны отображаться на набор "фискальных" устройств во время выполнения или, альтернативно, "устройство по умолчанию". Эта функция в настоящее время не существует. Добавление такой опции в ConfigProto - это то, что вы можете захотеть записать запрос функции.

Я могу предложить временное решение. Во-первых, некоторые комментарии к предлагаемым решениям.

  • Ваша первая идея, безусловно, будет работать, но, как вы указали, является громоздкой.

  • Настройка с использованием visible_device_list в ConfigProto не совсем работает, поскольку на самом деле это параметр для каждого процесса и игнорируется после создания первого сеанса в этом процессе. Это, конечно же, не документировано, как и должно быть (и несколько неудачно, что это появляется в настройке за сеанс). Однако это объясняет, почему ваше предложение здесь не работает и почему вы все еще видите один 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 для каждого графического процессора, используя что-то вроде:

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.

Надеюсь, что это поможет.

6
21 дек. '17 в 2:54
источник

Посмотрите другие вопросы по меткам или Задайте вопрос