java图形处理器_为多个模型指定CPU或GPU tensorflow java的工作

我正在使用Tensorflow java API(1.8.0),我加载了多个模型(在不同的会话中).使用SavedModelBundle.load(…)方法从.pb文件加载这些模型.这些.pb文件是通过保存Keras的模型获得的.

假设我要加载3个模型A,B,C.

为此,我实现了一个java Model类：

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);

}

}

然后,我可以轻松地使用此类实例化与我的A,B和C模型相对应的3个Model对象,并在同一个java程序中使用这3个模型进行预测.

我还注意到如果我有一个GPU,就会加载3个模型.

但是,我只希望模型A在GPU上运行并强制其他2个在CPU上运行.

通过阅读文档并深入了解源代码,我没有找到方法.我试图将可见设备的新ConfigProto定义为None,并使用图形实例化一个新的Session,但它不起作用(参见下面的代码).

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());

}

当我加载模型时,它使用可用的GPU.你对这个问题有什么解决方案吗？

谢谢您的回答.