Add auto layout to classification, detection and segmentation scripts

Signed-off-by: Serge Panev <spanev@nvidia.com>

Author: Kh4L

scripts/classification/imagenet/train_imagenet.py

@@ -108,6 +108,8 @@ def parse_args():
                         help='whether to use group norm.')
     parser.add_argument('--amp', action='store_true',
                     help='Use MXNet AMP for mixed precision training.')
+    parser.add_argument('--auto-layout', action='store_true',
+                        help='Add layout optimization to AMP. Must be used in addition of `--amp`.')
     opt = parser.parse_args()
     return opt
 
@@ -125,8 +127,10 @@ def main():
 
     logger.info(opt)
 
+    assert not opt.auto_layout or opt.amp, "--auto-layout needs to be used with --amp"
+
     if opt.amp:
-        amp.init()
+        amp.init(layout_optimization=opt.auto_layout)
 
     batch_size = opt.batch_size
     classes = 1000

scripts/detection/faster_rcnn/train_faster_rcnn.py

@@ -121,6 +121,8 @@ def parse_args():
                         help='Whether to use static memory allocation. Memory usage will increase.')
     parser.add_argument('--amp', action='store_true',
                         help='Use MXNet AMP for mixed precision training.')
+    parser.add_argument('--auto-layout', action='store_true',
+                        help='Add layout optimization to AMP. Must be used in addition of `--amp`.')
     parser.add_argument('--horovod', action='store_true',
                         help='Use MXNet Horovod for distributed training. Must be run with OpenMPI. '
                              '--gpus is ignored when using --horovod.')
@@ -622,8 +624,10 @@ def train(net, train_data, val_data, eval_metric, batch_size, ctx, args):
     # fix seed for mxnet, numpy and python builtin random generator.
     gutils.random.seed(args.seed)
 
+    assert not args.auto_layout or args.amp, "--auto-layout needs to be used with --amp"
+
     if args.amp:
-        amp.init()
+        amp.init(layout_optimization=args.auto_layout)
 
     # training contexts
     if args.horovod:

scripts/detection/ssd/train_ssd.py

@@ -89,6 +89,8 @@ def parse_args():
                         'Currently supports only COCO.')
     parser.add_argument('--amp', action='store_true',
                         help='Use MXNet AMP for mixed precision training.')
+    parser.add_argument('--auto-layout', action='store_true',
+                        help='Add layout optimization to AMP. Must be used in addition of `--amp`.')
     parser.add_argument('--horovod', action='store_true',
                         help='Use MXNet Horovod for distributed training. Must be run with OpenMPI. '
                         '--gpus is ignored when using --horovod.')
@@ -360,8 +362,10 @@ def train(net, train_data, val_data, eval_metric, ctx, args):
 if __name__ == '__main__':
     args = parse_args()
 
+    assert not args.auto_layout or args.amp, "--auto-layout needs to be used with --amp"
+
     if args.amp:
-        amp.init()
+        amp.init(layout_optimization=args.auto_layout)
 
     if args.horovod:
         hvd.init()

scripts/detection/yolo/train_yolo3.py

@@ -97,6 +97,8 @@ def parse_args():
     parser.add_argument('--label-smooth', action='store_true', help='Use label smoothing.')
     parser.add_argument('--amp', action='store_true',
                         help='Use MXNet AMP for mixed precision training.')
+    parser.add_argument('--auto-layout', action='store_true',
+                        help='Add layout optimization to AMP. Must be used in addition of `--amp`.')
     parser.add_argument('--horovod', action='store_true',
                         help='Use MXNet Horovod for distributed training. Must be run with OpenMPI. '
                         '--gpus is ignored when using --horovod.')
@@ -325,8 +327,10 @@ def train(net, train_data, val_data, eval_metric, ctx, args):
 if __name__ == '__main__':
     args = parse_args()
 
+    assert not args.auto_layout or args.amp, "--auto-layout needs to be used with --amp"
+
     if args.amp:
-        amp.init()
+        amp.init(layout_optimization=args.auto_layout)
 
     if args.horovod:
         if hvd is None:

scripts/instance/mask_rcnn/train_mask_rcnn.py

@@ -124,6 +124,8 @@ def parse_args():
                         help='Whether to use static memory allocation. Memory usage will increase.')
     parser.add_argument('--amp', action='store_true',
                         help='Use MXNet AMP for mixed precision training.')
+    parser.add_argument('--auto-layout', action='store_true',
+                        help='Add layout optimization to AMP. Must be used in addition of `--amp`.')
     parser.add_argument('--horovod', action='store_true',
                         help='Use MXNet Horovod for distributed training. Must be run with OpenMPI. '
                              '--gpus is ignored when using --horovod.')
@@ -700,8 +702,10 @@ def train(net, train_data, val_data, eval_metric, batch_size, ctx, logger, args)
     # fix seed for mxnet, numpy and python builtin random generator.
     gutils.random.seed(args.seed)
 
+    assert not args.auto_layout or args.amp, "--auto-layout needs to be used with --amp"
+
     if args.amp:
-        amp.init()
+        amp.init(layout_optimization=args.auto_layout)
 
     # training contexts
     if args.horovod:

scripts/segmentation/train.py

@@ -99,6 +99,8 @@ def parse_args():
     # performance related
     parser.add_argument('--amp', action='store_true',
                         help='Use MXNet AMP for mixed precision training.')
+    parser.add_argument('--auto-layout', action='store_true',
+                        help='Add layout optimization to AMP. Must be used in addition of `--amp`.')
     # handle contexts
     if args.no_cuda:
         print('Using CPU')
@@ -263,8 +265,10 @@ def save_checkpoint(net, args, epoch, mIoU, is_best=False):
 
 if __name__ == "__main__":
     args = parse_args()
+    assert not args.auto_layout or args.amp, "--auto-layout needs to be used with --amp"
+
     if args.amp:
-        amp.init()
+        amp.init(layout_optimization=args.auto_layout)
     # build logger
     filehandler = logging.FileHandler(os.path.join(args.save_dir, args.logging_file))
     streamhandler = logging.StreamHandler()