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

Question

我正在尝试使用掩码 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 中的值减少了一半。我还显着减少了图像和对象的数量。

还有其他想法吗？

Answer 1

500GB 是一个很好的内存量。我遇到了 GPU 内存不足的问题，这是一个单独的限制。
对于 TensorFlow v2，我发现以下有用:
1.减少batch_size到一个小的值
在配置文件中，设置:

train_config: {
  batch_size: 4
  ...
}

batch_size可以低至 1。
2. 减少调整大小的图像的尺寸
在配置文件中，设置 resizer height和 width到低于默认值 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))

这些变化中的每一个都有影响。但是，它们可能允许使用稀缺资源获得“足够好”的结果。