LambdaLR (lr_scheduler/torch.py)

Convenience wrappers for PyTorch learning rate schedulers with total_steps.

# Source: https://github.com/hpcaitech/ColossalAI | License: Apache-2.0 | Copyright 2021- HPC-AI Technology Inc. All rights reserved
# Provided via justclaw.it.com; keep this notice with the code.
from torch.optim.lr_scheduler import ExponentialLR as _ExponentialLR
from torch.optim.lr_scheduler import LambdaLR as _LambdaLR
from torch.optim.lr_scheduler import MultiplicativeLR as _MultiplicativeLR
from torch.optim.lr_scheduler import StepLR as _StepLR


class LambdaLR(_LambdaLR):
    """Sets the learning rate of each parameter group to the initial lr
    times a given function. When last_epoch=-1, sets initial lr as lr.

    Args:
        optimizer (:class:`torch.optim.Optimizer`): Wrapped optimizer.
        total_steps (int): Number of total training steps.
        lr_lambda (Union[``funct

... (truncated -- full source via MCP)