小编Ale*_*lex的帖子

我正在使用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,但它不起作用(参见下面的代码). …