Support specifying update interval in the parameter schedulers (#418)

Summary:
Pull Request resolved: #418

Most parameter schedulers used the default "epoch" update interval and didn't allow users to override this.
- This diff makes the update interval configurable by users
- The schedulers use a "step" interval by default, except for "step" and "multistep"
- Removed the default "epoch" update interval inside the base class - this isn't what most users would expect if they didn't specify the argument

Reviewed By: vreis

Differential Revision: D20256209

fbshipit-source-id: 1797d5d833dc20d173eed9ac53117e1e2de54ef1

Author: mannatsingh

classy_vision/optim/param_scheduler/classy_vision_param_scheduler.py

@@ -20,6 +20,25 @@ class UpdateInterval(Enum):
     EPOCH = "epoch"
     STEP = "step"
 
+    @classmethod
+    def from_config(
+        cls, config: Dict[str, Any], default: "UpdateInterval" = None
+    ) -> "UpdateInterval":
+        """Fetches the update interval from a config
+
+        Args:
+            config: The config for the parameter scheduler
+            default: The value to use if the config doesn't specify an update interval.
+                If not set, STEP is used.
+        """
+        if default is None:
+            default = cls.STEP
+        if "update_interval" not in config:
+            return default
+        if config.get("update_interval").lower() not in ["step", "epoch"]:
+            raise ValueError("Choices for update interval are 'step' or 'epoch'")
+        return cls[config["update_interval"].upper()]
+
 
 class ClassyParamScheduler(object):
     """
@@ -33,7 +52,7 @@ class ClassyParamScheduler(object):
     # To be used for comparisons with where
     WHERE_EPSILON = 1e-6
 
-    def __init__(self, update_interval: UpdateInterval = UpdateInterval.EPOCH):
+    def __init__(self, update_interval: UpdateInterval):
         """
         Constructor for ClassyParamScheduler
 

classy_vision/optim/param_scheduler/composite_scheduler.py

@@ -49,17 +49,17 @@ class CompositeParamScheduler(ClassyParamScheduler):
     The parameter value will be 0.42 for the first [0%, 30%) of steps,
     and then will cosine decay from 0.42 to 0.0001 for [30%, 100%) of
     training.
+    The schedule is updated after every train step by default.
     """
 
     def __init__(
         self,
         schedulers: Sequence[ClassyParamScheduler],
         lengths: Sequence[float],
-        update_interval: UpdateInterval,
         interval_scaling: Sequence[IntervalScaling],
+        update_interval: UpdateInterval = UpdateInterval.STEP,
     ):
-        super().__init__()
-        self.update_interval = update_interval
+        super().__init__(update_interval=update_interval)
         self._lengths = lengths
         self._schedulers = schedulers
         self._interval_scaling = interval_scaling
@@ -89,13 +89,6 @@ def from_config(cls, config: Dict[str, Any]) -> "CompositeParamScheduler":
         ), "The sum of all values in lengths must be 1"
         if sum(config["lengths"]) != 1.0:
             config["lengths"][-1] = 1.0 - sum(config["lengths"][:-1])
-        update_interval = UpdateInterval.STEP
-        if "update_interval" in config:
-            assert config["update_interval"] in {
-                "step",
-                "epoch",
-            }, "Choices for update interval are 'step' or 'epoch'"
-            update_interval = UpdateInterval[config["update_interval"].upper()]
         interval_scaling = []
         if "interval_scaling" in config:
             assert len(config["schedulers"]) == len(
@@ -119,7 +112,7 @@ def from_config(cls, config: Dict[str, Any]) -> "CompositeParamScheduler":
                 build_param_scheduler(scheduler) for scheduler in config["schedulers"]
             ],
             lengths=config["lengths"],
-            update_interval=update_interval,
+            update_interval=UpdateInterval.from_config(config, UpdateInterval.STEP),
             interval_scaling=interval_scaling,
         )
 

classy_vision/optim/param_scheduler/constant_scheduler.py

@@ -6,7 +6,7 @@
 
 from typing import Any, Dict
 
-from . import ClassyParamScheduler, register_param_scheduler
+from . import ClassyParamScheduler, UpdateInterval, register_param_scheduler
 
 
 @register_param_scheduler("constant")
@@ -16,7 +16,7 @@ class ConstantParamScheduler(ClassyParamScheduler):
     """
 
     def __init__(self, value: float):
-        super().__init__()
+        super().__init__(update_interval=UpdateInterval.EPOCH)
         self._value = value
 
     @classmethod

classy_vision/optim/param_scheduler/cosine_scheduler.py

@@ -7,7 +7,7 @@
 import math
 from typing import Any, Dict
 
-from . import ClassyParamScheduler, register_param_scheduler
+from . import ClassyParamScheduler, UpdateInterval, register_param_scheduler
 
 
 @register_param_scheduler("cosine")
@@ -17,6 +17,7 @@ class CosineParamScheduler(ClassyParamScheduler):
     //arxiv.org/pdf/1608.03983.pdf>`_.
     Can be used for either cosine decay or cosine warmup schedules based on
     start and end values.
+    The schedule is updated after every train step by default.
 
     Example:
 
@@ -26,8 +27,13 @@ class CosineParamScheduler(ClassyParamScheduler):
           end_value: 0.0001
     """
 
-    def __init__(self, start_value: float, end_value: float):
-        super().__init__()
+    def __init__(
+        self,
+        start_value: float,
+        end_value: float,
+        update_interval: UpdateInterval = UpdateInterval.STEP,
+    ):
+        super().__init__(update_interval=update_interval)
         self._start_value = start_value
         self._end_value = end_value
 
@@ -46,7 +52,11 @@ def from_config(cls, config: Dict[str, Any]) -> "CosineParamScheduler":
             "start_value" in config and "end_value" in config
         ), "Cosine scheduler requires a start_value and a end_value"
 
-        return cls(start_value=config["start_value"], end_value=config["end_value"])
+        return cls(
+            start_value=config["start_value"],
+            end_value=config["end_value"],
+            update_interval=UpdateInterval.from_config(config, UpdateInterval.STEP),
+        )
 
     def __call__(self, where: float):
         return self._end_value + 0.5 * (self._start_value - self._end_value) * (

classy_vision/optim/param_scheduler/linear_scheduler.py

@@ -6,14 +6,15 @@
 
 from typing import Any, Dict
 
-from . import ClassyParamScheduler, register_param_scheduler
+from . import ClassyParamScheduler, UpdateInterval, register_param_scheduler
 
 
 @register_param_scheduler("linear")
 class LinearParamScheduler(ClassyParamScheduler):
     """
     Linearly interpolates parameter between ``start_value`` and ``end_value``.
     Can be used for either warmup or decay based on start and end values.
+    The schedule is updated after every train step by default.
 
     Example:
 
@@ -24,8 +25,13 @@ class LinearParamScheduler(ClassyParamScheduler):
     Corresponds to a linear increasing schedule with values in [0.0001, 0.01)
     """
 
-    def __init__(self, start_value: float, end_value: float):
-        super().__init__()
+    def __init__(
+        self,
+        start_value: float,
+        end_value: float,
+        update_interval: UpdateInterval = UpdateInterval.STEP,
+    ):
+        super().__init__(update_interval=update_interval)
         self._start_value = start_value
         self._end_value = end_value
 
@@ -43,7 +49,12 @@ def from_config(cls, config: Dict[str, Any]) -> "LinearParamScheduler":
         assert (
             "start_value" in config and "end_value" in config
         ), "Linear scheduler requires a start and a end"
-        return cls(start_value=config["start_value"], end_value=config["end_value"])
+
+        return cls(
+            start_value=config["start_value"],
+            end_value=config["end_value"],
+            update_interval=UpdateInterval.from_config(config, UpdateInterval.STEP),
+        )
 
     def __call__(self, where: float):
         # interpolate between start and end values

classy_vision/optim/param_scheduler/multi_step_scheduler.py

@@ -18,6 +18,7 @@ class MultiStepParamScheduler(ClassyParamScheduler):
     """
     Takes a predefined schedule for a param value, and a list of epochs
     which stand for the upper boundary (excluded) of each range.
+    The schedule is updated after every train epoch by default.
 
     Example:
 
@@ -37,10 +38,10 @@ def __init__(
         self,
         values,
         num_epochs: int,
-        update_interval: UpdateInterval,
         milestones: Optional[List[int]] = None,
+        update_interval: UpdateInterval = UpdateInterval.EPOCH,
     ):
-        super().__init__(update_interval)
+        super().__init__(update_interval=update_interval)
         self._param_schedule = values
         self._num_epochs = num_epochs
         self._milestones = milestones
@@ -96,11 +97,12 @@ def from_config(cls, config: Dict[str, Any]) -> "MultiStepParamScheduler":
                 "Non-Equi Step scheduler requires a list of %d epochs"
                 % (len(config["values"]) - 1)
             )
+
         return cls(
             values=config["values"],
             num_epochs=config["num_epochs"],
-            update_interval=UpdateInterval(config.get("update_interval", "epoch")),
             milestones=milestones,
+            update_interval=UpdateInterval.from_config(config, UpdateInterval.EPOCH),
         )
 
     def __call__(self, where: float):

classy_vision/optim/param_scheduler/polynomial_decay_scheduler.py

@@ -6,14 +6,15 @@
 
 from typing import Any, Dict
 
-from . import ClassyParamScheduler, register_param_scheduler
+from . import ClassyParamScheduler, UpdateInterval, register_param_scheduler
 
 
 @register_param_scheduler("polynomial")
 class PolynomialDecayParamScheduler(ClassyParamScheduler):
     """
     Decays the param value after every epoch according to a
     polynomial function with a fixed power.
+    The schedule is updated after every train step by default.
 
     Example:
 
@@ -26,8 +27,13 @@ class PolynomialDecayParamScheduler(ClassyParamScheduler):
     so on.
     """
 
-    def __init__(self, base_value, power):
-        super().__init__()
+    def __init__(
+        self,
+        base_value: float,
+        power: float,
+        update_interval: UpdateInterval = UpdateInterval.STEP,
+    ):
+        super().__init__(update_interval=update_interval)
 
         self._base_value = base_value
         self._power = power
@@ -46,7 +52,11 @@ def from_config(cls, config: Dict[str, Any]) -> "PolynomialDecayParamScheduler":
         assert (
             "base_value" in config and "power" in config
         ), "Polynomial decay scheduler requires a base lr and a power of decay"
-        return cls(base_value=config["base_value"], power=config["power"])
+        return cls(
+            base_value=config["base_value"],
+            power=config["power"],
+            update_interval=UpdateInterval.from_config(config, UpdateInterval.STEP),
+        )
 
     def __call__(self, where: float):
         return self._base_value * (1 - where) ** self._power

classy_vision/optim/param_scheduler/step_scheduler.py

@@ -6,7 +6,7 @@
 
 from typing import Any, Dict, List, NamedTuple, Optional, Union
 
-from . import ClassyParamScheduler, register_param_scheduler
+from . import ClassyParamScheduler, UpdateInterval, register_param_scheduler
 
 
 @register_param_scheduler("step")
@@ -15,6 +15,7 @@ class StepParamScheduler(ClassyParamScheduler):
     Takes a fixed schedule for a param value.  If the length of the
     fixed schedule is less than the number of epochs, then the epochs
     are divided evenly among the param schedule.
+    The schedule is updated after every train epoch by default.
 
     Example:
 
@@ -27,8 +28,13 @@ class StepParamScheduler(ClassyParamScheduler):
     epochs 30-59, 0.001 for epoch 60-89, 0.0001 for epochs 90-119.
     """
 
-    def __init__(self, num_epochs: Union[int, float], values: List[float]):
-        super().__init__()
+    def __init__(
+        self,
+        num_epochs: Union[int, float],
+        values: List[float],
+        update_interval: UpdateInterval = UpdateInterval.EPOCH,
+    ):
+        super().__init__(update_interval=update_interval)
 
         self._param_schedule = values
 
@@ -50,7 +56,11 @@ def from_config(cls, config: Dict[str, Any]) -> "StepParamScheduler":
         ), "Step scheduler requires a list of at least one param value"
         assert config["num_epochs"] > 0, "Num epochs must be greater than 0"
 
-        return cls(num_epochs=config["num_epochs"], values=config["values"])
+        return cls(
+            num_epochs=config["num_epochs"],
+            values=config["values"],
+            update_interval=UpdateInterval.from_config(config, UpdateInterval.EPOCH),
+        )
 
     def __call__(self, where: float):
         ind = int((where + self.WHERE_EPSILON) * len(self._param_schedule))

classy_vision/optim/param_scheduler/step_with_fixed_gamma_scheduler.py

@@ -15,6 +15,7 @@ class StepWithFixedGammaParamScheduler(ClassyParamScheduler):
     """
     Decays the param value by gamma at equal number of steps so as to have the
     specified total number of decays.
+    The schedule is updated after every train step by default.
 
     Example:
 
@@ -29,6 +30,31 @@ class StepWithFixedGammaParamScheduler(ClassyParamScheduler):
     epochs 30-59, 0.001 for epoch 60-89, 0.0001 for epochs 90-119.
     """
 
+    def __init__(
+        self,
+        base_value: float,
+        num_decays: int,
+        gamma: float,
+        num_epochs: int,
+        update_interval: UpdateInterval = UpdateInterval.STEP,
+    ):
+        super().__init__(update_interval=update_interval)
+
+        self.base_value = base_value
+        self.num_decays = num_decays
+        self.gamma = gamma
+        self.num_epochs = num_epochs
+        values = [base_value]
+        for _ in range(num_decays):
+            values.append(values[-1] * gamma)
+
+        self._step_param_scheduler = StepParamScheduler(
+            num_epochs=num_epochs, values=values
+        )
+
+        # make this a STEP scheduler
+        self.update_interval = UpdateInterval.STEP
+
     @classmethod
     def from_config(cls, config: Dict[str, Any]) -> "StepWithFixedGammaParamScheduler":
         """Instantiates a StepWithFixedGammaParamScheduler from a configuration.
@@ -56,25 +82,8 @@ def from_config(cls, config: Dict[str, Any]) -> "StepWithFixedGammaParamSchedule
             num_decays=config["num_decays"],
             gamma=config["gamma"],
             num_epochs=config["num_epochs"],
+            update_interval=UpdateInterval.from_config(config, UpdateInterval.STEP),
         )
 
-    def __init__(self, base_value, num_decays, gamma, num_epochs):
-        super().__init__()
-
-        self.base_value = base_value
-        self.num_decays = num_decays
-        self.gamma = gamma
-        self.num_epochs = num_epochs
-        values = [base_value]
-        for _ in range(num_decays):
-            values.append(values[-1] * gamma)
-
-        self._step_param_scheduler = StepParamScheduler(
-            num_epochs=num_epochs, values=values
-        )
-
-        # make this a STEP scheduler
-        self.update_interval = UpdateInterval.STEP
-
     def __call__(self, where: float) -> float:
         return self._step_param_scheduler(where)

test/optim_param_scheduler_composite_test.py

@@ -93,7 +93,7 @@ def test_invalid_config(self):
         # Bad value for update_interval
         bad_config["lengths"] = copy.deepcopy(config["lengths"])
         bad_config["update_interval"] = "epochs"
-        with self.assertRaises(AssertionError):
+        with self.assertRaises(Exception):
             CompositeParamScheduler.from_config(bad_config)
 
         # Bad value for composition_mode