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

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

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

Ограничить видимость устройства на уровне процесса, а именно с помощью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.

Создайте модель в многоуровневом режиме на Python, используяtf.device() Спецификация. На java стороне дают разныеPredictors различные башни внутри общего сеанса.

Мне тяжеловато и идиоматически неправильно.

ОБНОВИТЬ: Как предложил @ash, есть еще один вариант:

Назначьте соответствующее устройство для каждой операции существующего графа, изменив его определение (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 tenorflow? Любые идеи были бы хорошы.