Fix a logical error of mask rcnn training script (#1249)

chinakook · web-flow · commit db8559f44095 · 2020-04-24T19:14:01.000-07:00
* Fix a logical error of mask rcnn training script

After this fix, the training memory will be steady. If given a very big epoch&#65288;often with small dataset), this error will eat all GPU memory as epoch grows.
The script of faster rcnn training has this problem too. I&lsquo;ll fix it later when I get some tests passed.

* fix hybridize
diff --git a/scripts/detection/faster_rcnn/train_faster_rcnn.py b/scripts/detection/faster_rcnn/train_faster_rcnn.py
@@ -539,13 +539,13 @@ def train(net, train_data, val_data, eval_metric, batch_size, ctx, args):
         logger.info(net.collect_train_params().keys())
     logger.info('Start training from [Epoch {}]'.format(args.start_epoch))
     best_map = [0]
+    rcnn_task = ForwardBackwardTask(net, trainer, rpn_cls_loss, rpn_box_loss, rcnn_cls_loss,
+                                    rcnn_box_loss, mix_ratio=1.0)
+    executor = Parallel(args.executor_threads, rcnn_task) if not args.horovod else None
     for epoch in range(args.start_epoch, args.epochs):
         mix_ratio = 1.0
         if not args.disable_hybridization:
             net.hybridize(static_alloc=args.static_alloc)
-        rcnn_task = ForwardBackwardTask(net, trainer, rpn_cls_loss, rpn_box_loss, rcnn_cls_loss,
-                                        rcnn_box_loss, mix_ratio=1.0)
-        executor = Parallel(args.executor_threads, rcnn_task) if not args.horovod else None
         if args.mixup:
             # TODO(zhreshold) only support evenly mixup now, target generator needs to be modified otherwise
             train_data._dataset._data.set_mixup(np.random.uniform, 0.5, 0.5)
diff --git a/scripts/instance/mask_rcnn/train_mask_rcnn.py b/scripts/instance/mask_rcnn/train_mask_rcnn.py
@@ -615,12 +615,12 @@ def train(net, train_data, val_data, eval_metric, batch_size, ctx, logger, args)
     logger.info('Start training from [Epoch {}]'.format(args.start_epoch))
     best_map = [0]
     base_lr = trainer.learning_rate
+    rcnn_task = ForwardBackwardTask(net, trainer, rpn_cls_loss, rpn_box_loss, rcnn_cls_loss,
+                                    rcnn_box_loss, rcnn_mask_loss)
+    executor = Parallel(args.executor_threads, rcnn_task) if not args.horovod else None
     for epoch in range(args.start_epoch, args.epochs):
         if not args.disable_hybridization:
             net.hybridize(static_alloc=args.static_alloc)
-        rcnn_task = ForwardBackwardTask(net, trainer, rpn_cls_loss, rpn_box_loss, rcnn_cls_loss,
-                                        rcnn_box_loss, rcnn_mask_loss)
-        executor = Parallel(args.executor_threads, rcnn_task) if not args.horovod else None
         while lr_steps and epoch >= lr_steps[0]:
             new_lr = trainer.learning_rate * lr_decay
             lr_steps.pop(0)
