Support specifying end lr for WarmupCosineLR.

Summary:
Support specifying end lr for WarmupCosineLR
* Use `cfg.SOLVER.BASE_LR_END` to specify the lr value for the last iteration, only used by `WarmupCosineLR`.

Reviewed By: zhanghang1989

Differential Revision: D33292501

fbshipit-source-id: 07942d734f8d445fe03f82a85bdc78b65a8aebd0

Author: newstzpz

detectron2/config/defaults.py

@@ -517,6 +517,8 @@
 _C.SOLVER.MAX_ITER = 40000
 
 _C.SOLVER.BASE_LR = 0.001
+# The end lr, only used by WarmupCosineLR
+_C.SOLVER.BASE_LR_END = 0.0
 
 _C.SOLVER.MOMENTUM = 0.9
 

detectron2/solver/build.py

@@ -272,7 +272,9 @@ def build_lr_scheduler(
             num_updates=cfg.SOLVER.MAX_ITER,
         )
     elif name == "WarmupCosineLR":
-        sched = CosineParamScheduler(1, 0)
+        end_value = cfg.SOLVER.BASE_LR_END / cfg.SOLVER.BASE_LR
+        assert end_value >= 0.0 and end_value <= 1.0, end_value
+        sched = CosineParamScheduler(1, end_value)
     else:
         raise ValueError("Unknown LR scheduler: {}".format(name))
 

tests/test_scheduler.py

@@ -1,14 +1,14 @@
 # Copyright (c) Facebook, Inc. and its affiliates.
 
 import math
-import numpy as np
 from unittest import TestCase
+
+import numpy as np
 import torch
+from detectron2.solver import LRMultiplier, WarmupParamScheduler, build_lr_scheduler
 from fvcore.common.param_scheduler import CosineParamScheduler, MultiStepParamScheduler
 from torch import nn
 
-from detectron2.solver import LRMultiplier, WarmupParamScheduler
-
 
 class TestScheduler(TestCase):
     def test_warmup_multistep(self):
@@ -66,3 +66,50 @@ def test_warmup_cosine(self):
                 self.assertAlmostEqual(lr, expected_cosine)
             else:
                 self.assertNotAlmostEqual(lr, expected_cosine)
+
+    def test_warmup_cosine_end_value(self):
+        from detectron2.config import CfgNode, get_cfg
+
+        def _test_end_value(cfg_dict):
+            cfg = get_cfg()
+            cfg.merge_from_other_cfg(CfgNode(cfg_dict))
+
+            p = nn.Parameter(torch.zeros(0))
+            opt = torch.optim.SGD([p], lr=cfg.SOLVER.BASE_LR)
+
+            scheduler = build_lr_scheduler(cfg, opt)
+
+            p.sum().backward()
+            opt.step()
+            self.assertEqual(
+                opt.param_groups[0]["lr"], cfg.SOLVER.BASE_LR * cfg.SOLVER.WARMUP_FACTOR
+            )
+
+            lrs = []
+            for _ in range(cfg.SOLVER.MAX_ITER):
+                scheduler.step()
+                lrs.append(opt.param_groups[0]["lr"])
+
+            self.assertAlmostEqual(lrs[-1], cfg.SOLVER.BASE_LR_END)
+
+        _test_end_value({
+            "SOLVER": {
+                "LR_SCHEDULER_NAME": "WarmupCosineLR",
+                "MAX_ITER": 100,
+                "WARMUP_ITERS": 10,
+                "WARMUP_FACTOR": 0.1,
+                "BASE_LR": 5.0,
+                "BASE_LR_END": 0.0,
+            }
+        })
+
+        _test_end_value({
+            "SOLVER": {
+                "LR_SCHEDULER_NAME": "WarmupCosineLR",
+                "MAX_ITER": 100,
+                "WARMUP_ITERS": 10,
+                "WARMUP_FACTOR": 0.1,
+                "BASE_LR": 5.0,
+                "BASE_LR_END": 0.5,
+            }
+        })