Укажите процессор или графический процессор для нескольких моделей tenorflow java
Я использую Tensorflow java API (1.8.0), где загружаю несколько моделей (в разных сеансах). Эти модели загружаются из файлов .pb с использованием метода SavedModelBundle.load (...). Эти файлы .pb были получены путем сохранения моделей Кераса.
Допустим, я хочу загрузить 3 модели A, B, C. Для этого я реализовал класс java Модель:
public class Model implements Closeable {
private String inputName;
private String outputName;
private Session session;
private int inputSize;
public Model(String modelDir, String input_name, String output_name, int inputSize) {
SavedModelBundle bundle = SavedModelBundle.load(modelDir, "serve");
this.inputName = input_name;
this.outputName = output_name;
this.inputSize = inputSize;
this.session = bundle.session();
}
public void close() {
session.close();
}
public Tensor predict(Tensor t) {
return session.runner().feed(inputName, t).fetch(outputName).run().get(0);
}
}
Затем я легко могу создать экземпляры 3 объектов Модель, соответствующих моим моделям A, B и C, с этим классом и сделать прогнозы с этими 3 моделями в той же программе Java. Я также заметил, что если у меня есть графический процессор, на него загружаются 3 модели.
Однако Я хотел бы, чтобы только модель A работала на графическом процессоре, а две другие - на CPU.
Прочитав документацию и погрузившись в исходный код, я не нашел способа сделать это. Я попытался определить новую настройку ConfigProto для видимых устройств на None и создать экземпляр нового сеанса с графиком, но это не сработало (см. Код ниже).
public Model(String modelDir, String input_name, String output_name, int inputSize) {
SavedModelBundle bundle = SavedModelBundle.load(modelDir, "serve");
this.inputName = input_name;
this.outputName = output_name;
this.inputSize = inputSize;
ConfigProto configProto = ConfigProto.newBuilder().setAllowSoftPlacement(false).setGpuOptions(GPUOptions.newBuilder().setVisibleDeviceList("").build()).build();
this.session = new Session(bundle.graph(),configProto.toByteArray());
}
Когда я загружаю модель, она использует доступный графический процессор. У вас есть решение этой проблемы?
Спасибо за ваш ответ.
Ответы 3
Вы можете установить конфигурацию device вашего графика тензорного потока. Вот соответствующий код: [источник].
...
byte[] config = ConfigProto.newBuilder()
.setLogDevicePlacement(true)
.setAllowSoftPlacement(true)
.build()
.toByteArray()
Session sessions[] = new Session[numModels];
// function to move the graph definition to a new device
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();
}
graphA.importGraphDef(modifyGraphDef(graphDef, "/gpu:0"));
graphB.importGraphDef(modifyGraphDef(graphDef, "/cpu:0"));
Это, вероятно, было бы чище, чем более очевидная установка переменной среды CUDA_VISIBLE_DEVICES на "" после загрузки первой модели.
@alex в качестве некоторой отладки, устанавливает CUDA_VISIBLE_DEVICES в пустую строку, предотвращает переход всех моделей на GPU (как и должно)
Да, это так, но использовать это решение для меня невозможно, так как для этого требуется изменить переменную среды.
Согласно Эта проблема, новый исходный код устранил эту проблему. К сожалению, вам придется собирать из исходников после эти инструкции
Затем вы можете протестировать:
ConfigProto configProto = ConfigProto.newBuilder()
.setAllowSoftPlacement(true) // allow less GPUs than configured
.setGpuOptions(GPUOptions.newBuilder().setPerProcessGpuMemoryFraction(0.01).build())
.build();
SavedModelBundle bundle = SavedModelBundle.loader(modelDir).withTags("serve").withConfigProto(configProto.toByteArray()).load();
Приведенные выше ответы не помогли мне. Использование putDeviceCount ("GPU", 0) заставляет TF использовать CPU. Он работает в версии 1.15.0. Вы можете загружать одну и ту же модель как в процессор, так и в графический процессор, и если графический процессор выдает Resource is made: OOM при выделении тензора, используйте модель процессора для прогнозирования.
ConfigProto configProtoCpu = ConfigProto.newBuilder().setAllowSoftPlacement(true).putDeviceCount("GPU", 0)
.build();
SavedModelBundle modelCpu=SavedModelBundle.loader(modelPath).withTags("serve")
.withConfigProto(configProtoCpu.toByteArray()).load();
ConfigProto configProtoGpu = ConfigProto.newBuilder().setAllowSoftPlacement(true)
.setGpuOptions(GPUOptions.newBuilder().setAllowGrowth(true).build()).build();
SavedModelBundle modelgpu = SavedModelBundle.loader(modelPath).withTags("serve")
.withConfigProto(configProtoGpu.toByteArray()).load();
Я не знаком с Java API Tensorflow. Но я знаю, что с его API Python вы можете делать что-то вроде:
with tf.device('/cpu:0'): # your graphилиwith tf.device('/gpu:0'): # your graph. Я думаю, что в Java должно быть что-то подобное. Итак, я поискал и нашел этот ответ на SO. Я думаю (особенно последний фрагмент кода) - это решение, которое вы ищете, но я не уверен. Пожалуйста, подтвердите это.