dual_scheduler steps may happen multiple times per epoch

[juan43ramirez]

Bug

The call to dual_scheduler.step() happens at the end of ConstrainedOptimizer.step(). This means that dual_scheduler steps happen at the end of every optimization step and not every optimization epoch.

https://github.com/cooper-org/cooper/blob/eae7c5a68563b83410913c8b24e99b70e32694f3/cooper/constrained_optimizer.py#L235-L240

This goes against the Pytorch convention on learning rate schedulers.

Expected behavior

Calls to dual_scheduler.step() should happen at the end of every epoch.

How to fix

We could: (i) Remove the call to dual_scheduler.step() from inside the ConstrainedOptimizer and let the user handle it. The instantiation of the dual scheduler could still be performed internally so that the partial_lr_scheduler logic is preserved (and for the user to not worry about it). (ii) Indicate whether a step is the last of its epoch and only then call dual_scheduler.step(). (iii) Create a class which wraps both the primal and dual schedulers (like how the ConstrainedOptimizer wraps primal and dual optimizers), whose step() method is manually called by the user at the end of an epoch. The class could internally handle the initialization of the dual scheduler. This approach is similar to (i)

[juan43ramirez]

This bug has been fixed by #36. The current version of the code expects the user to perform the call to constrained_optimizer.dual_scheduler.step() explicitly. This way the user can ensure that the changes on the dual scheduler align with the end of an epoch. See the documentation on dual schedulers for more details.