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

【问题标题】：Specify either CPU or GPU for multiple models tensorflow java's job为多个模型 tensorflow java 的作业指定 CPU 或 GPU
【发布时间】：2018-11-22 00:20:10
【问题描述】：

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

感谢您的回答。

【问题讨论】：

我对Tensorflow的Java API不熟悉。但我知道使用它的 Python API，您可以执行以下操作：with tf.device('/cpu:0'): # your graph 或 with tf.device('/gpu:0'): # your graph。我认为它在 Java 中也必须是类似的东西。所以我在 SO 上搜索并找到了this answer。我认为（特别是最后一段代码）是您正在寻找的解决方案，但我不确定。请确认。

标签： java tensorflow keras gpu

【解决方案1】：

您可以设置张量流图的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 环境变量更明显地设置为"" 更简洁。

【讨论】：

@alex 作为一些调试，是否将 CUDA_VISIBLE_DEVICES 设置为空字符串会阻止所有模型在 GPU 上运行（应该如此）
是的，可以，但我无法使用此解决方案，因为它需要修改环境变量。

【解决方案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();

【讨论】：

【解决方案3】：

上面给出的答案对我不起作用。使用 putDeviceCount("GPU", 0) 使 TF 使用 CPU 。它在 1.15.0 版本中运行。您可以将相同的模型加载到 cpu 和 gpu，如果 gpu 抛出 Resource exhausted: 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();

【讨论】：