Question

我正在使用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。你有解决这个问题的方法吗？

感谢您的回答。

Answer 1

您可以设置张量流图的date配置。这是一些相关的代码accepted answer。

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更干净。

Answer 2

根据this issue，新的源代码解决了此问题。不幸的是，您必须遵循these instructions

从源代码构建

然后您可以测试：

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

Answer 3

以上给出的答案对我不起作用。使用putDeviceCount（“ GPU”，0）使TF使用CPU。它正在1.15.0版中运行。您可以将相同的模型加载到cpu和gpu上，并且如果gpu抛出资源耗尽：分配张量时为OOM，请使用CPU模型进行预测。

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

为多个模型指定CPU或GPU tensorflow java的工作

3 个答案: