feat: adagrad integration

metaopt · Benjamin-eecs · Mar 22, 2023 · Oct 13, 2022 · Feb 15, 2023 · Feb 15, 2023
commit ace6b5efa517c38b4b88737d08d2a9d349ac2f56
diff --git a/Makefile b/Makefile
@@ -114,7  114,7 @@ addlicense-install: go-install
 
 pytest: test-install
 	cd tests && $(PYTHON) -c 'import $(PROJECT_NAME)' && \
-	$(PYTHON) -m pytest -k "test_exponential_decay" --verbose --color=yes --durations=0 \
 	$(PYTHON) -m pytest --verbose --color=yes --durations=0 \
 		--cov="$(PROJECT_NAME)" --cov-config=.coveragerc --cov-report=xml --cov-report=term-missing \
 		$(PYTESTOPTS) .
 

diff --git a/torchopt/alias/adagrad.py b/torchopt/alias/adagrad.py
@@ -39,16  39,15 @@
     scale_by_neg_lr,
 )
 from torchopt.combine import chain
-from torchopt.transform import scale_by_rss
-from torchopt.typing import GradientTransformation, Numeric, Scalar, Schedule
 from torchopt.transform import scale_by_rss, scale_by_schedule
 from torchopt.typing import GradientTransformation, Numeric, Scalar, ScalarOrSchedule, Schedule
 
 
 __all__ = ['adagrad']
 
 
 # pylint: disable-next=too-many-arguments
-def _adagrad_lr_decay(
-    init_value: Scalar,
 def _adagrad_lr_schedule(
     decay_rate: Scalar,
     transition_begin: int = 0,
 ) -> Schedule:
@@ -61,10  60,10 @@ def _adagrad_lr_decay(
     ```
 
     Args:
-        init_value: the initial learning rate.
-        decay_rate: The decay rate.
-        transition_begin: must be positive. After how many steps to start annealing
         decay_rate (float, optional): The decay rate.
         transition_begin (int, optional): must be positive. After how many steps to start annealing
             (before this many steps the scalar value is held fixed at `init_value`).
             (default: :const:`1`)
 
     Returns:
         schedule: A function that maps step counts to values.
@@ -78,14  77,14 @@ def _adagrad_lr_decay(
 
     def schedule(count: Numeric) -> Numeric:
         decreased_count = count - transition_begin
-        return init_value / (1   decay_rate * decreased_count)
         return 1 / (1   decay_rate * decreased_count)
 
     return schedule
 
 
 # pylint: disable-next=too-many-arguments
 def adagrad(
-    lr: Scalar = 1e-2,
     lr: ScalarOrSchedule = 1e-2,
     lr_decay: float = 0.0,
     weight_decay: float = 0.0,
     initial_accumulator_value: float = 0.0,
@@ -144,21  143,22 @@ def adagrad(
     flip_sign_and_add_weight_decay_fn = flip_sign_and_add_weight_decay
     adagrad_scaler_fn = scale_by_rss
     scale_by_neg_lr_fn = scale_by_neg_lr
-    schedule_fn = _adagrad_lr_decay
     step_size_fn = _adagrad_lr_schedule
     scale_by_schedule_fn = scale_by_schedule
 
     if _get_use_chain_flat():  # default behavior
         chain_fn = chain_fn.flat  # type: ignore[attr-defined]
         flip_sign_and_add_weight_decay_fn = flip_sign_and_add_weight_decay_fn.flat  # type: ignore[attr-defined]
         adagrad_scaler_fn = adagrad_scaler_fn.flat  # type: ignore[attr-defined]
         scale_by_neg_lr_fn = scale_by_neg_lr_fn.flat  # type: ignore[attr-defined]
         scale_by_schedule_fn = scale_by_schedule_fn.flat  # type: ignore[attr-defined]
 
     return chain_fn(
         flip_sign_and_add_weight_decay_fn(weight_decay=weight_decay, maximize=maximize),
         adagrad_scaler_fn(
             initial_accumulator_value=initial_accumulator_value,
             eps=eps,
         ),
-        scale_by_neg_lr_fn(
-            schedule_fn(init_value=lr, decay_rate=lr_decay, transition_begin=0),
-        ),
         scale_by_schedule_fn(step_size_fn=step_size_fn(decay_rate=lr_decay, transition_begin=0)),
         scale_by_neg_lr_fn(lr),
     )
diff --git a/torchopt/schedule/exponential_decay.py b/torchopt/schedule/exponential_decay.py
@@ -71,7  71,7 @@ def exponential_decay(
             entire annealing process is disabled and the value is held fixed at ``init_value``.
             (default: :const:`1`)
         staircase (bool): If ``True``, decay the scalar at discrete intervals.
-        end_value (float or Tensor): End value of the scalar to be annealed.
         end_value (float or Tensor, optional): End value of the scalar to be annealed.
 
     Returns:
         schedule: A function that maps step counts to values.