Add AMP to ImageNet segmentation script

Signed-off-by: Serge Panev <[email protected]>

Author: Kh4L

scripts/segmentation/train.py

@@ -8,6 +8,7 @@
 import mxnet as mx
 from mxnet import gluon, autograd
 from mxnet.gluon.data.vision import transforms
+from mxnet.contrib import amp
 
 import gluoncv
 gluoncv.utils.check_version('0.6.0')
@@ -95,7 +96,9 @@ def parse_args():
                         help='using Synchronized Cross-GPU BatchNorm')
     # the parser
     args = parser.parse_args()
-
+    # performance related
+    parser.add_argument('--amp', action='store_true',
+                        help='Use MXNet AMP for mixed precision training.')
     # handle contexts
     if args.no_cuda:
         print('Using CPU')
@@ -200,7 +203,11 @@ def __init__(self, args, logger):
                 v.wd_mult = 0.0
 
         self.optimizer = gluon.Trainer(self.net.module.collect_params(), 'sgd',
-                                       optimizer_params, kvstore=kv)
+                                       optimizer_params, kvstore=(False if args.amp else None))
+
+
+        if args.amp:
+            amp.init_trainer(trainer)
         # evaluation metrics
         self.metric = gluoncv.utils.metrics.SegmentationMetric(trainset.num_class)
 
@@ -212,7 +219,11 @@ def training(self, epoch):
                 outputs = self.net(data.astype(args.dtype, copy=False))
                 losses = self.criterion(outputs, target)
                 mx.nd.waitall()
-                autograd.backward(losses)
+                if args.amp:
+                    with amp.scale_loss(losses, self.optimizer) as scaled_losses:
+                        autograd.backward(scaled_losses)
+                else:
+                    autograd.backward(losses)
             self.optimizer.step(self.args.batch_size)
             for loss in losses:
                 train_loss += np.mean(loss.asnumpy()) / len(losses)
@@ -252,7 +263,8 @@ def save_checkpoint(net, args, epoch, mIoU, is_best=False):
 
 if __name__ == "__main__":
     args = parse_args()
-
+    if args.amp:
+        amp.init()
     # build logger
     filehandler = logging.FileHandler(os.path.join(args.save_dir, args.logging_file))
     streamhandler = logging.StreamHandler()