Tensorflow对象检测掩码rcnn使用了太多内存

时间:2018-03-03 05:08:50

标签: tensorflow

我正在尝试使用掩码rcnn运行TF对象检测,但它会在具有500GB内存的节点上继续死亡。

我将models / research / object_detection / trainer.py ConfigProto更新为

session_config = tf.ConfigProto(allow_soft_placement=True,
                                intra_op_parallelism_threads=1,
                                inter_op_parallelism_threads=1,
                                device_count = {'CPU': 1},
                                log_device_placement=False)

我将mask_rcnn_inception_resnet_v2_atrous_coco.config更新为

train_config: {
  batch_queue_capacity: 500
  num_batch_queue_threads: 8
  prefetch_queue_capacity: 10

到目前为止,更新ConfigProto效果最佳。在它死亡之前我已经完成了30步而不是1.我将train_config中的值减少了一半。我也大大减少了图像和物体的数量。

还有其他想法吗?

3 个答案:

答案 0 :(得分:1)

我有类似的问题。通过设置以下值,我设法将内存消耗减少了2.5倍:

prefetch_size: 4
num_readers: 4
min_after_dequeue: 1

我不确定他们中的哪一个(可能全部?)负责减少内存,(我没有测试过)或者他们的确切值会影响内存消耗,但你可以轻松地尝试一下。

答案 1 :(得分:0)

一些以前用来减少内存使用的选项已被弃用。来自object_detection/protos/input_reader.proto

optional uint32 queue_capacity = 3 [default=2000, deprecated=true];
optional uint32 min_after_dequeue = 4 [default=1000, deprecated=true];
optional uint32 prefetch_size = 13 [default = 512, deprecated=true];
optional uint32 num_parallel_map_calls = 14 [default = 64, deprecated=true];

到目前为止, num_parallel_batches 似乎是最大的存储猪。

我的配置文件中的*_input_reader消息现在看起来像这样:

train_input_reader: {
  tf_record_input_reader {
    input_path: "<DATASET_DIR>/tfrecords/train*.tfrecord"
  }
  label_map_path: "<DATASET_DIR>/label_map.pbtxt"
  load_instance_masks: true
  mask_type: PNG_MASKS
  num_parallel_batches: 1
}

Mask RCNN训练现在使用的CPU内存比以前减少了约50%(在775 x 522图像上进行训练)。

答案 2 :(得分:0)

500GB是足够的内存。我遇到了GPU内存不足的问题,这是一个单独的约束。

对于TensorFlow v2,我发现以下有用:

1。将batch_size减小到较小的值

在配置文件中,设置:

train_config: {
  batch_size: 4
  ...
}

batch_size可以低至1。

2。缩小尺寸调整后的图像的尺寸

在配置文件中,将大小调整器heightwidth设置为小于默认值1024x1024的值。

model {
  faster_rcnn {
    number_of_stages: 3
    num_classes: 1
    image_resizer {
      fixed_shape_resizer {
        height: 256
        width: 256
      }
    }

3。不要训练功能检测器

这仅适用于Mask R-CNN,并且是最难实施的更改。在文件research/object_detection/model_lib_v2.py中,更改以下代码:

当前:

def eager_train_step(detection_model,
...
  trainable_variables = detection_model.trainable_variables
  gradients = tape.gradient(total_loss, trainable_variables)

  if clip_gradients_value:
    gradients, _ = tf.clip_by_global_norm(gradients, clip_gradients_value)
  optimizer.apply_gradients(zip(gradients, trainable_variables))

新功能:

def eager_train_step(detection_model,
...
  # Mask R-CNN variables to train -- not feature detector
  trainable_variables = detection_model.trainable_variables
  to_fine_tune = []
  prefixes_to_train = ['FirstStageBoxPredictor',
                       'mask_rcnn_keras_box_predictor',
                       'RPNConv'
                        ]
  for var in trainable_variables:
    if any([var.name.startswith(prefix) for prefix in prefixes_to_train]):
      to_fine_tune.append(var)
    
  gradients = tape.gradient(total_loss, to_fine_tune)

  if clip_gradients_value:
    gradients, _ = tf.clip_by_global_norm(gradients, clip_gradients_value)
  optimizer.apply_gradients(zip(gradients, to_fine_tune))

每个更改都有其含义。但是,它们可以使用稀缺资源获得“足够好”的结果。