Method __init__

cdslib/core/optim/tf_optim.py:15–109

r""" Vary the learning rate every training step according to: .. math:: lr = \text{factor} * d_{model}^{-0.5} * \min(\text{iter}^{-0.5}, \text{iter} * \text{warmup}^{-1.5}) This corresponds to linearly increasing the learning rate for the first `warmup`

(
        self,
        optimizer: torch.optim.Optimizer,
        model_size: int = 512,
        factor: float = 1.0,
        warmup: int = 2000,
        init_step: int = 0,
        use_dynamic_rate: bool = True,
    )

    """

    def __init__(
        self,
        optimizer: torch.optim.Optimizer,
        model_size: int = 512,
        factor: float = 1.0,
        warmup: int = 2000,
        init_step: int = 0,
        use_dynamic_rate: bool = True,
    ):
        r"""
        Vary the learning rate every training step according to:

            .. math::
                lr = \text{factor} * d_{model}^{-0.5} * \min(\text{iter}^{-0.5}, \text{iter} * \text{warmup}^{-1.5})

        This corresponds to linearly increasing the learning rate for the first `warmup` training steps,
        and after that decreasing it proportionally to the inverse square root of the iteration (step).
        This optimizer was used by Vaswani, Ashish, et al. ("Attention is all you need." 2017).


        Args:
            optimizer (torch.optim.Optimizer):
                pytorch optimizer to warp around
            model_size (int):
                primary dimension of the model
            factor (float):
                scalar to multiply. It is to scale the learning rate.
            warmup (int):
                number of warm up steps to linearly increase the learning rate with
            init_step (int):
                initial step (for training recovery)
            use_dynamic_rate (bool):
                whether to (True) use the increasing-decreasing schedule
                or (False) the constant learning rate

        Examples:
            .. code-block:: python

                use_amp = True  # whether to use Automatic Mixed Precision
                net = SomeModel(...)
                parameters = net.parameters()

                # Create a pytorch optimizer
                # (it will handle everything other than learning rate,
                # e.g., normalization, momentum, etc)
                _optimizer = torch.optim.Adam(
                        parameters,  # parameters to optimize
                        lr=1e-3, # this will be overwritten by our optimzier
                        betas=(0.9, 0.98),
                        eps=1e-9)

                # wrap our learning rate scheduler
                optimizer = TFOptimizer(
                        optimizer=_optimizer,
                        model_size=512,
                        factor=1.0,
                        warmup=4000,
                        init_step=0)