I am attempting to make a custom scheduler which replicates the Roberta warmup, so far I came up with this, based on Huggingface’s linear warmup scheduler

  
      github.com
  
  
    huggingface/transformers/blob/master/src/transformers/optimization.py#L71 4

    """    def lr_lambda(current_step: int):        if current_step < num_warmup_steps:            return float(current_step) / float(max(1.0, num_warmup_steps))        return 1.0    return LambdaLR(optimizer, lr_lambda, last_epoch=last_epoch)def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):    """    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after    a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.    Args:        optimizer (:class:`~torch.optim.Optimizer`):            The optimizer for which to schedule the learning rate.        num_warmup_steps (:obj:`int`):            The number of steps for the warmup phase.        num_training_steps (:obj:`int`):



  
  
    
    
  
  


def get_linear_schedule_with_warmup_with_peak(optimizer, num_warmup_steps, num_training_steps, init_lr, peak_lr, last_epoch=-1):

    def lr_lambda(current_step: int):
        if current_step < num_warmup_steps:
            return (float(current_step) / float(max(1, num_warmup_steps)))*(peak_lr/init_lr)
        return max(
            0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))
        )

    return LambdaLR(optimizer, lr_lambda, last_epoch)

Which is not quite exact, since the optimizer is AdamW, which has its own weight decay. For it to be exact, I need init_lr to be replaced by the current optimizer learning rate, but from the documentation on LambdaLR, lr_lambda seems to only take in an int.

  
      pytorch.org
  
  
    

torch.optim — PyTorch 1.7.0 documentation 5



  
  
    
    
  
  


So I need a way to adjust the optimizer’s learning rate based on its current learning rate.