add a convient scheduler just for warmup

Summary: it's a special case of CompositeParamScheduler

Reviewed By: theschnitz

Differential Revision: D26246409

fbshipit-source-id: afb5d49d99c8be237c59a464c27184763b4db150

Author: ppwwyyxx

detectron2/solver/__init__.py

@@ -1,5 +1,5 @@
 # Copyright (c) Facebook, Inc. and its affiliates.
 from .build import build_lr_scheduler, build_optimizer, get_default_optimizer_params
-from .lr_scheduler import WarmupCosineLR, WarmupMultiStepLR, LRMultiplier
+from .lr_scheduler import WarmupCosineLR, WarmupMultiStepLR, LRMultiplier, WarmupParamScheduler
 
 __all__ = [k for k in globals().keys() if not k.startswith("_")]

detectron2/solver/build.py

@@ -3,17 +3,11 @@
 from enum import Enum
 from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Type, Union
 import torch
-from fvcore.common.param_scheduler import (
-    CompositeParamScheduler,
-    ConstantParamScheduler,
-    CosineParamScheduler,
-    LinearParamScheduler,
-    MultiStepParamScheduler,
-)
+from fvcore.common.param_scheduler import CosineParamScheduler, MultiStepParamScheduler
 
 from detectron2.config import CfgNode
 
-from .lr_scheduler import LRMultiplier
+from .lr_scheduler import LRMultiplier, WarmupParamScheduler
 
 _GradientClipperInput = Union[torch.Tensor, Iterable[torch.Tensor]]
 _GradientClipper = Callable[[_GradientClipperInput], None]
@@ -221,18 +215,10 @@ def build_lr_scheduler(
     else:
         raise ValueError("Unknown LR scheduler: {}".format(name))
 
-    # Add warmup
-    warmup_method = cfg.SOLVER.WARMUP_METHOD
-    if warmup_method == "constant":
-        warmup = ConstantParamScheduler(cfg.SOLVER.WARMUP_FACTOR)
-    elif warmup_method == "linear":
-        warmup = LinearParamScheduler(cfg.SOLVER.WARMUP_FACTOR, 1.0)
-    else:
-        raise ValueError("Unknown warmup method: {}".format(warmup_method))
-    warmup_ratio = cfg.SOLVER.WARMUP_ITERS / cfg.SOLVER.MAX_ITER
-    sched = CompositeParamScheduler(
-        [warmup, sched],
-        interval_scaling=["rescaled", "fixed"],
-        lengths=[warmup_ratio, 1 - warmup_ratio],
+    sched = WarmupParamScheduler(
+        sched,
+        cfg.SOLVER.WARMUP_FACTOR,
+        cfg.SOLVER.WARMUP_ITERS / cfg.SOLVER.MAX_ITER,
+        cfg.SOLVER.WARMUP_METHOD,
     )
     return LRMultiplier(optimizer, multiplier=sched, max_iter=cfg.SOLVER.MAX_ITER)

detectron2/solver/lr_scheduler.py

@@ -4,11 +4,51 @@
 from bisect import bisect_right
 from typing import List
 import torch
-from fvcore.common.param_scheduler import ParamScheduler
+from fvcore.common.param_scheduler import (
+    CompositeParamScheduler,
+    ConstantParamScheduler,
+    LinearParamScheduler,
+    ParamScheduler,
+)
 
 logger = logging.getLogger(__name__)
 
 
+class WarmupParamScheduler(CompositeParamScheduler):
+    """
+    Add an initial warmup stage to another scheduler.
+    """
+
+    def __init__(
+        self,
+        scheduler: ParamScheduler,
+        warmup_factor: float,
+        warmup_length: float,
+        warmup_method: str = "linear",
+    ):
+        """
+        Args:
+            scheduler: warmup will be added at the beginning of this scheduler
+            warmup_factor: the factor w.r.t the initial value of ``scheduler``, e.g. 0.001
+            warmup_length: the relative length (in [0, 1]) of warmup steps w.r.t the entire
+                training, e.g. 0.01
+            warmup_method: one of "linear" or "constant"
+        """
+        end_value = scheduler(warmup_length)  # the value to reach when warmup ends
+        start_value = warmup_factor * scheduler(0.0)
+        if warmup_method == "constant":
+            warmup = ConstantParamScheduler(start_value)
+        elif warmup_method == "linear":
+            warmup = LinearParamScheduler(start_value, end_value)
+        else:
+            raise ValueError("Unknown warmup method: {}".format(warmup_method))
+        super().__init__(
+            [warmup, scheduler],
+            interval_scaling=["rescaled", "fixed"],
+            lengths=[warmup_length, 1 - warmup_length],
+        )
+
+
 class LRMultiplier(torch.optim.lr_scheduler._LRScheduler):
     """
     A LRScheduler which uses fvcore :class:`ParamScheduler` to multiply the
@@ -25,15 +65,12 @@ class LRMultiplier(torch.optim.lr_scheduler._LRScheduler):
     ::
         LRMultiplier(
             opt,
-            CompositeParamScheduler([
-                LinearParamScheduler(0.001, 1),  # warmup
+            WarmupParamScheduler(
                 MultiStepParamScheduler(
                     [1, 0.1, 0.01],
                     milestones=[60000, 80000],
                     num_updates=90000,
-                )],
-                interval_scaling=["rescaled", "fixed"],
-                lengths=[100 / 90000, 89900 / 90000],
+                ), 0.001, 100 / 90000
             ),
             max_iter=90000
         )

tests/test_scheduler.py

@@ -4,33 +4,25 @@
 import numpy as np
 from unittest import TestCase
 import torch
-from fvcore.common.param_scheduler import (
-    CompositeParamScheduler,
-    CosineParamScheduler,
-    LinearParamScheduler,
-    MultiStepParamScheduler,
-)
+from fvcore.common.param_scheduler import CosineParamScheduler, MultiStepParamScheduler
 from torch import nn
 
-from detectron2.solver import LRMultiplier
+from detectron2.solver import LRMultiplier, WarmupParamScheduler
 
 
 class TestScheduler(TestCase):
     def test_warmup_multistep(self):
         p = nn.Parameter(torch.zeros(0))
         opt = torch.optim.SGD([p], lr=5)
 
-        multiplier = CompositeParamScheduler(
-            [
-                LinearParamScheduler(0.001, 1),  # warmup
-                MultiStepParamScheduler(
-                    [1, 0.1, 0.01, 0.001],
-                    milestones=[10, 15, 20],
-                    num_updates=30,
-                ),
-            ],
-            interval_scaling=["rescaled", "fixed"],
-            lengths=[5 / 30, 25 / 30],
+        multiplier = WarmupParamScheduler(
+            MultiStepParamScheduler(
+                [1, 0.1, 0.01, 0.001],
+                milestones=[10, 15, 20],
+                num_updates=30,
+            ),
+            0.001,
+            5 / 30,
         )
         sched = LRMultiplier(opt, multiplier, 30)
         # This is an equivalent of:
@@ -53,13 +45,10 @@ def test_warmup_multistep(self):
     def test_warmup_cosine(self):
         p = nn.Parameter(torch.zeros(0))
         opt = torch.optim.SGD([p], lr=5)
-        multiplier = CompositeParamScheduler(
-            [
-                LinearParamScheduler(0.001, 1),  # warmup
-                CosineParamScheduler(1, 0),
-            ],
-            interval_scaling=["rescaled", "fixed"],
-            lengths=[5 / 30, 25 / 30],
+        multiplier = WarmupParamScheduler(
+            CosineParamScheduler(1, 0),
+            0.001,
+            5 / 30,
         )
         sched = LRMultiplier(opt, multiplier, 30)