def cosine_warmup_schedule(
    step: int,
    max_steps: int,
    start_value: float,
    end_value: float,
    warmup_steps: int,
    warmup_start_value: float,
    warmup_end_value: Optional[float] = None,
    period: Optional[int] = None,
) -> float:
    """Use cosine decay to gradually modify start_value to reach target end_value.

    Uses linear warmup for the first warmup_steps steps.

    Args:
        step:
            Current step number.
        max_steps:
            Total number of steps.
        start_value:
            Starting value.
        end_value:
            Target value.
        warmup_steps:
            Number of steps for warmup.
        warmup_start_value:
            Starting value for warmup.
        warmup_end_value:
            Target value for warmup. Defaults to start_value.
        period:
            The number of steps over which the cosine function completes a full cycle.
            If no period is provided, the scheduler will complete a half cycle over
            max_steps - warmup_steps.
    Returns:
        Cosine decay value.
    """
    if warmup_steps < 0:
        raise ValueError(f"Warmup steps {warmup_steps} can't be negative.")
    if warmup_steps > max_steps:
        raise ValueError(f"Warmup steps {warmup_steps} must be <= max_steps.")
    if step > max_steps:
        warnings.warn(
            f"Current step number {step} exceeds max_steps {max_steps}.",
            category=RuntimeWarning,
        )

    if warmup_end_value is None:
        warmup_end_value = start_value

    if step < warmup_steps:
        # Use step + 1 to reach warmup_end_value at end of warmup. This means that the
        # initial warmup_start_value is skipped which is oftentimes desired when setting
        # it to 0 as this would result in no parameter updates.
        return (
            warmup_start_value
            + (warmup_end_value - warmup_start_value) * (step + 1) / warmup_steps
        )
    else:
        max_steps = max_steps - (warmup_steps if period is None else 1)
        return cosine_schedule(
            step=step - warmup_steps,
            max_steps=max_steps,
            start_value=start_value,
            end_value=end_value,
            period=period,
        )
Use cosine decay to gradually modify start_value to reach target end_value.

    Uses linear warmup for the first warmup_steps steps.

    Args:
        step:
            Current step number.
        max_steps:
            Total number of steps.
        start_value:
            Starting value.
        end_value:
            Target value.
        warmup_steps:
            Number of steps for warmup.
        warmup_start_value:
            Starting value for warmup.
        warmup_end_value:
            Target value for warmup. Defaults to start_value.
        period:
            The number of steps over which the cosine function completes a full cycle.
            If no period is provided, the scheduler will complete a half cycle over
            max_steps - warmup_steps.
    Returns:
        Cosine decay value.
    
lightly/utils/scheduler.py