Tensorflow Java多GPU推理

时间：2017-12-13 18:35:03

标签： java tensorflow tensorflow-gpu multi-gpu

我有一台带有多个GPU的服务器，并希望在Java应用程序内的模型推理期间充分利用它们。 默认情况下，tensorflow会占用所有可用的GPU，但只使用第一个GPU。

我可以想出三个选项来克服这个问题：

  1. 限制设备在流程级别上的可见性，即使用CUDA_VISIBLE_DEVICES环境变量。

    这将要求我运行java应用程序的几个实例并在它们之间分配流量。不是那个诱人的想法。

  2. 在单个应用程序中启动多个会话，并尝试通过ConfigProto为每个会话分配一个设备：

    public class DistributedPredictor {

    private Predictor[] nested;
    private int[] counters;

    // ...

    public DistributedPredictor(String modelPath, int numDevices, int numThreadsPerDevice) {
        nested = new Predictor[numDevices];
        counters = new int[numDevices];

        for (int i = 0; i < nested.length; i++) {
            nested[i] = new Predictor(modelPath, i, numDevices, numThreadsPerDevice);
        }
    }

    public Prediction predict(Data data) {
        int i = acquirePredictorIndex();
        Prediction result = nested[i].predict(data);
        releasePredictorIndex(i);
        return result;
    }

    private synchronized int acquirePredictorIndex() {
        int i = argmin(counters);
        counters[i] += 1;
        return i;
    }

    private synchronized void releasePredictorIndex(int i) {
        counters[i] -= 1;
    }
}


public class Predictor {

    private Session session;

    public Predictor(String modelPath, int deviceIdx, int numDevices, int numThreadsPerDevice) {

        GPUOptions gpuOptions = GPUOptions.newBuilder()
                .setVisibleDeviceList("" + deviceIdx)
                .setAllowGrowth(true)
                .build();

        ConfigProto config = ConfigProto.newBuilder()
                .setGpuOptions(gpuOptions)
                .setInterOpParallelismThreads(numDevices * numThreadsPerDevice)
                .build();

        byte[] graphDef = Files.readAllBytes(Paths.get(modelPath));
        Graph graph = new Graph();
        graph.importGraphDef(graphDef);

        this.session = new Session(graph, config.toByteArray());
    }

    public Prediction predict(Data data) {
        // ...
    }
}

    这种方法似乎一目了然。但是，会话偶尔会忽略setVisibleDeviceList选项，并且所有会话都会导致第一个设备导致内存崩溃。

  3. 使用tf.device()规范在python中以多塔方式构建模型。在java方面，在共享会话中提供不同的Predictor个不同的塔。

    对我来说感到麻烦和惯用错误。

    4. 更新：正如@ash提议的那样，还有另一种选择：

    1. 通过修改现有图形的定义（graphDef），为现有图形的每个操作分配一个合适的设备。

      要完成它，可以调整方法2中的代码：

      public class Predictor {

    private Session session;

    public Predictor(String modelPath, int deviceIdx, int numDevices, int numThreadsPerDevice) {

        byte[] graphDef = Files.readAllBytes(Paths.get(modelPath));
        graphDef = setGraphDefDevice(graphDef, deviceIdx)

        Graph graph = new Graph();
        graph.importGraphDef(graphDef);

        ConfigProto config = ConfigProto.newBuilder()
                .setAllowSoftPlacement(true)
                .build();

        this.session = new Session(graph, config.toByteArray());
    }

    private static byte[] setGraphDefDevice(byte[] graphDef, int deviceIdx) throws InvalidProtocolBufferException {
        String deviceString = String.format("/gpu:%d", deviceIdx);

        GraphDef.Builder builder = GraphDef.parseFrom(graphDef).toBuilder();
        for (int i = 0; i < builder.getNodeCount(); i++) {
            builder.getNodeBuilder(i).setDevice(deviceString);
        }
        return builder.build().toByteArray();
    }

    public Prediction predict(Data data) {
        // ...
    }
}

      就像其他提到的方法一样，这个方法并没有让我免于在设备之间手动分配数据。但至少它运行稳定，并且相对容易实现。总的来说，这看起来像（几乎）正常的技术。

      2. 使用tensorflow java API有一种优雅的方式来做这样的基本操作吗？任何想法，将不胜感激。

2 个答案:

答案 0 :(得分：6)

简而言之：有一种解决方法，每个GPU最终会有一个会话。

详细说明：

一般流程是TensorFlow运行时尊重为图中的操作指定的设备。如果没有为某个操作指定任何设备，那么它就会放置＆＃34;它基于一些启发式方法。这些启发式算法目前导致GPU上的操作：0如果GPU可用且存在用于操作的GPU内核＆＃34; （Placer::Run如果您有兴趣）。

我要求的是我认为对TensorFlow的合理功能请求 - 将序列化图形中的设备视为＆＃34;虚拟＆＃34;要映射到一组＆＃34; phyiscal＆＃34;设备在运行时，或者设置＆＃34;默认设备＆＃34;。此功能目前不存在。将此类选项添加到ConfigProto是您可能要为其提交功能请求的内容。

我可以在此期间建议解决方法。首先，对您提出的解决方案进行一些评论。

  1. 你的第一个想法肯定会奏效，但正如你所指出的那样，很麻烦。

  2. 使用visible_device_list中的ConfigProto进行设置并不能完全解决，因为这实际上是一个按进程设置，并且在进程中创建第一个会话后会被忽略。这肯定没有记录，应该是（并且有点不幸的是，它出现在每会话配置中）。但是，这解释了为什么您的建议不起作用以及您仍然看到使用单个GPU的原因。

  3. 这可行。

    4. 另一种选择是最终得到不同的图形（操作明确放在不同的GPU上），每个GPU产生一个会话。这样的东西可以用来编辑图形并明确地为每个操作分配一个设备：

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

    之后，您可以使用以下内容为每个GPU创建Graph和Session 

    final int NUM_GPUS = 8;
// setAllowSoftPlacement: Just in case our device modifications were too aggressive
// (e.g., setting a GPU device on an operation that only has CPU kernels)
// setLogDevicePlacment: So we can see what happens.
byte[] config =
    ConfigProto.newBuilder()
        .setLogDevicePlacement(true)
        .setAllowSoftPlacement(true)
        .build()
        .toByteArray();
Graph graphs[] = new Graph[NUM_GPUS];
Session sessions[] = new Session[NUM_GPUS];
for (int i = 0; i < NUM_GPUS; ++i) {
  graphs[i] = new Graph();
  graphs[i].importGraphDef(modifyGraphDef(graphDef, String.format("/gpu:%d", i)));
  sessions[i] = new Session(graphs[i], config);    
}

    然后使用sessions[i]在GPU #i上执行图表。

    希望有所帮助。

答案 1 :(得分：0)

在python中，可以按照以下步骤进行操作：

def get_frozen_graph(graph_file):
    """Read Frozen Graph file from disk."""
    with tf.gfile.GFile(graph_file, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    return graph_def

trt_graph1 = get_frozen_graph('/home/ved/ved_1/frozen_inference_graph.pb')

with tf.device('/gpu:1'):
    [tf_input_l1, tf_scores_l1, tf_boxes_l1, tf_classes_l1, tf_num_detections_l1, tf_masks_l1] = tf.import_graph_def(trt_graph1, 
                    return_elements=['image_tensor:0', 'detection_scores:0', 
                    'detection_boxes:0', 'detection_classes:0','num_detections:0', 'detection_masks:0'])
    
tf_sess1 = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))

trt_graph2 = get_frozen_graph('/home/ved/ved_2/frozen_inference_graph.pb')

with tf.device('/gpu:0'):
    [tf_input_l2, tf_scores_l2, tf_boxes_l2, tf_classes_l2, tf_num_detections_l2] = tf.import_graph_def(trt_graph2, 
                    return_elements=['image_tensor:0', 'detection_scores:0', 
                    'detection_boxes:0', 'detection_classes:0','num_detections:0'])
    
tf_sess2 = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))
相关问题
最新问题