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

Question

我正在使用Tensorflow java API（1.8.0），在这里我加载多个模型（在不同的会话中）。这些模型是从加载的。pb使用保存的模型绑定的文件。负载（…） 方法。那些。pb文件是通过保存Keras的模型获得的

假设我想加载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);
}
}

然后我可以很容易地用这个类实例化与我的A、B和C模型对应的3个模型对象，并在同一个java程序中用这3个模型进行预测。 我还注意到，如果我有一个GPU，3个模型都加载在上面

然而，我只希望模型A在GPU上运行，并强制其他2个在CPU上运行

通过阅读文档和深入研究源代码，我没有找到这样做的方法。我试图定义一个新的ConfigProto，将visible devices设置为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());
}

当我加载模型时，它使用可用的GPU。你有解决这个问题的办法吗

谢谢你的回答

Answer 1

根据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 2

您可以设置tensorflow图的device配置。下面是一些相关的代码 [source]

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