Tensorflow Java多GPU推理

Question

我有一台配备了多个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在共享会话中分配不同的塔。

   对我来说感觉笨重且用法不正确。

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

4. 通过修改其定义(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) {
           // ...
       }
   }
   

   就像其他提到的方法一样，这种方法并不能使我从手动分配数据给不同设备中获得解放。但至少它工作稳定，实现起来相对容易。总的来说，这看起来像是一种（几乎）正常的技术。

在tensorflow java API中有一种优雅的方法来做这样基本的事情吗？欢迎任何想法。

- Alexander Lutsenko

Answer 1

简而言之：有一种解决方法，它可以让每个GPU都拥有一个会话。

详情：

TensorFlow runtime将尊重图中指定操作的设备。如果没有为操作指定设备，则根据某些启发式规则"放置"它。这些启发式规则目前会导致"如果GPU可用并且有GPU内核，则在GPU:0上放置操作"（如果您感兴趣，可以查看Placer::Run）。

我认为您要求的是TensorFlow的一个合理特性请求，即将序列化图中的设备视为“虚拟”设备，在运行时映射到一组“物理”设备，或者设置“默认设备”。此功能目前不存在。将这种选项添加到ConfigProto中是您可能想提出特性请求的事情。

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

1. 您的第一个想法肯定有效，但正如您指出的那样，它很麻烦。 2. 在ConfigProto中使用 visible_device_list 不太奏效，因为它实际上是一个进程级别的设置，并且在进程创建第一个会话后被忽略。这显然没有得到很好的文档说明（并且有点令人遗憾，因为它出现在每个Session的配置中）。但是，这解释了为什么您的建议不起作用，以及为什么您仍然只看到使用了一个GPU。 3. 这可以行得通。

另一种选择是得到不同的图（将操作显式放置在不同的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();
} 

接下来，您可以创建一个Graph和Session，每个GPU都使用类似于以下内容的东西：

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上执行图形。

希望这可以帮到你。

- ash

Answer 2

在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))

- Vedanshu