Decrease the maximum learning rate after every restart
I'm training a neural network for a computer vision-based task. For the optimizer, I found out that it isn't ideal to use a single learning rate for the entire training, and what people do is that they use learning rate schedulers to decay the learning rate in a specific manner. So to do this, I tried out PyTorch's CosineAnnealingWarmRestarts(). What this does is that it anneals/decreases the initial learning rate (set by us) in a cosine manner until it hits a restart. After this "restart," the learning rate is set back to the initial learning rate, and the cycle happens again. This worked pretty well for me, but I wanted to make a few changes in it. I wanted to change the learning rate, the optimizer is assigned after each restart, so that after every restart the maximum learning rate for the optimizer also decreases. Can this be done in PyTorch?

It seems to me that a straight-forward solution would just be to inherit from CosineAnnealingWarmRestarts and then change its self.optimizer parameters inside an overriden step function. In pseudo-code, that would be something like

class myScheduler(torch.optim.lr_scheduler.CosineAnnealingWarmRestarts):
    def __init__(self,
            optimizer, 
            T_0, 
            T_mult=1,
            eta_min=0,
            last_epoch=-1):
    #initialize base calss
    super().__init__(.... blablabla ...)

    def step(self):
        #call step() from base class
        #Do some book-keeping to determine if you've hit a restart 
        #now change optimizer lr for each parameter group
        if some_condition:#condition like number of iterations, restarts, etc
            self.optimizer.param_groups[i]['lr']*=some_coef

