fix(detection): fix learning rate schedule (#114)

Author: Jianfeng Wang

official/vision/classification/resnet/train.py

@@ -165,11 +165,11 @@ def valid_step(image, label):
 
     # multi-step learning rate scheduler with warmup
     def adjust_learning_rate(step):
-        lr = args.lr * 0.1 ** bisect.bisect_right(
+        lr = args.lr * dist.get_world_size() * 0.1 ** bisect.bisect_right(
             [30 * steps_per_epoch, 60 * steps_per_epoch, 80 * steps_per_epoch], step
         )
         if step < 5 * steps_per_epoch:  # warmup
-            lr = args.lr * (step / (5 * steps_per_epoch))
+            lr = args.lr * dist.get_world_size() * (step / (5 * steps_per_epoch))
         for param_group in opt.param_groups:
             param_group["lr"] = lr
         return lr

official/vision/classification/shufflenet/train.py

@@ -176,7 +176,7 @@ def valid_step(image, label):
 
     # linear learning rate scheduler
     def adjust_learning_rate(step):
-        lr = args.lr * (1 - step / (args.epochs * steps_per_epoch))
+        lr = args.lr * dist.get_world_size() * (1 - step / (args.epochs * steps_per_epoch))
         for param_group in opt.param_groups:
             param_group["lr"] = lr
         return lr

official/vision/detection/README.md

@@ -43,7 +43,7 @@
 
 ## 安装和环境配置
 
-本目录下代码基于MegEngine v1.2，在开始运行本目录下的代码之前，请确保按照[README](../../../README.md)进行了正确的环境配置。
+本目录下代码基于MegEngine v1.6，在开始运行本目录下的代码之前，请确保按照[README](../../../README.md)进行了正确的环境配置。
 
 ## 如何使用
 

official/vision/detection/tools/test.py

@@ -20,7 +20,6 @@
 
 logger = mge.get_logger(__name__)
 logger.setLevel("INFO")
-mge.device.set_prealloc_config(1024, 1024, 256 * 1024 * 1024, 4.0)
 
 
 def make_parser():

official/vision/detection/tools/test_random.py

@@ -24,7 +24,6 @@
 
 logger = mge.get_logger(__name__)
 logger.setLevel("INFO")
-mge.device.set_prealloc_config(1024, 1024, 256 * 1024 * 1024, 4.0)
 
 
 def make_parser():

official/vision/detection/tools/train.py

@@ -30,7 +30,6 @@
 
 logger = mge.get_logger(__name__)
 logger.setLevel("INFO")
-mge.device.set_prealloc_config(1024, 1024, 256 * 1024 * 1024, 4.0)
 
 
 def make_parser():
@@ -183,7 +182,7 @@ def train_func(image, im_info, gt_boxes):
 
 def adjust_learning_rate(optimizer, epoch, step, cfg, args):
     base_lr = (
-        cfg.basic_lr * args.batch_size * (
+        cfg.basic_lr * args.batch_size * dist.get_world_size() * (
             cfg.lr_decay_rate
             ** bisect.bisect_right(cfg.lr_decay_stages, epoch)
         )

official/vision/detection/tools/train_random.py

@@ -29,7 +29,6 @@
 
 logger = mge.get_logger(__name__)
 logger.setLevel("INFO")
-mge.device.set_prealloc_config(1024, 1024, 256 * 1024 * 1024, 4.0)
 
 
 def make_parser():
@@ -182,7 +181,7 @@ def train_func(image, im_info, gt_boxes):
 
 def adjust_learning_rate(optimizer, epoch, step, cfg, args):
     base_lr = (
-        cfg.basic_lr * args.batch_size * (
+        cfg.basic_lr * args.batch_size * dist.get_world_size() * (
             cfg.lr_decay_rate
             ** bisect.bisect_right(cfg.lr_decay_stages, epoch)
         )