chore: import upstream snapshot with attribution

labml_nn/optimizers/configs.py

@@ -0,0 +1,156 @@
+"""
+---
+title: Configurable optimizer module
+summary: This implements a configurable module for optimizers.
+---
+
+# Configurable Optimizer
+"""
+
+from typing import Tuple
+
+import torch
+
+from labml.configs import BaseConfigs, option, meta_config
+from labml_nn.optimizers import WeightDecay
+
+
+class OptimizerConfigs(BaseConfigs):
+    """
+    <a id="OptimizerConfigs"></a>
+
+    ## Optimizer Configurations
+    """
+
+    # Optimizer
+    optimizer: torch.optim.Adam
+
+    # Weight decay
+    weight_decay_obj: WeightDecay
+    # Whether weight decay is decoupled;
+    # i.e. weight decay is not added to gradients
+    weight_decouple: bool = True
+    # Weight decay
+    weight_decay: float = 0.0
+    # Whether weight decay is absolute or should be multiplied by learning rate
+    weight_decay_absolute: bool = False
+
+    # Whether the adam update is optimized (different epsilon)
+    optimized_adam_update: bool = True
+
+    # Parameters to be optimized
+    parameters: any
+
+    # Learning rate $\alpha$
+    learning_rate: float = 0.01
+    # Beta values $(\beta_1, \beta_2)$ for Adam
+    betas: Tuple[float, float] = (0.9, 0.999)
+    # Epsilon $\epsilon$ for adam
+    eps: float = 1e-08
+
+    # Momentum for SGD
+    momentum: float = 0.5
+    # Whether to use AMSGrad
+    amsgrad: bool = False
+
+    # Number of warmup optimizer steps
+    warmup: int = 2_000
+    # Total number of optimizer steps (for cosine decay)
+    total_steps: int = int(1e10)
+
+    # Whether to degenerate to SGD in AdaBelief
+    degenerate_to_sgd: bool = True
+
+    # Whether to use Rectified Adam in AdaBelief
+    rectify: bool = True
+
+    # Model embedding size for Noam optimizer
+    d_model: int
+
+    rho: float
+
+    def __init__(self):
+        super().__init__(_primary='optimizer')
+
+
+meta_config(OptimizerConfigs.parameters)
+
+
+@option(OptimizerConfigs.weight_decay_obj, 'L2')
+def _weight_decay(c: OptimizerConfigs):
+    return WeightDecay(c.weight_decay, c.weight_decouple, c.weight_decay_absolute)
+
+
+@option(OptimizerConfigs.optimizer, 'SGD')
+def _sgd_optimizer(c: OptimizerConfigs):
+    return torch.optim.SGD(c.parameters, c.learning_rate, c.momentum,
+                           weight_decay=c.weight_decay)
+
+
+@option(OptimizerConfigs.optimizer, 'Adam')
+def _adam_optimizer(c: OptimizerConfigs):
+    if c.amsgrad:
+        from labml_nn.optimizers.amsgrad import AMSGrad
+        return AMSGrad(c.parameters,
+                       lr=c.learning_rate, betas=c.betas, eps=c.eps,
+                       optimized_update=c.optimized_adam_update,
+                       weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad)
+    else:
+        from labml_nn.optimizers.adam import Adam
+        return Adam(c.parameters,
+                    lr=c.learning_rate, betas=c.betas, eps=c.eps,
+                    optimized_update=c.optimized_adam_update,
+                    weight_decay=c.weight_decay_obj)
+
+
+@option(OptimizerConfigs.optimizer, 'AdamW')
+def _adam_warmup_optimizer(c: OptimizerConfigs):
+    from labml_nn.optimizers.adam_warmup import AdamWarmup
+    return AdamWarmup(c.parameters,
+                      lr=c.learning_rate, betas=c.betas, eps=c.eps,
+                      weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad, warmup=c.warmup)
+
+
+@option(OptimizerConfigs.optimizer, 'RAdam')
+def _radam_optimizer(c: OptimizerConfigs):
+    from labml_nn.optimizers.radam import RAdam
+    return RAdam(c.parameters,
+                 lr=c.learning_rate, betas=c.betas, eps=c.eps,
+                 weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad,
+                 degenerated_to_sgd=c.degenerate_to_sgd)
+
+
+@option(OptimizerConfigs.optimizer, 'AdaBelief')
+def _ada_belief_optimizer(c: OptimizerConfigs):
+    from labml_nn.optimizers.ada_belief import AdaBelief
+    return AdaBelief(c.parameters,
+                     lr=c.learning_rate, betas=c.betas, eps=c.eps,
+                     weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad,
+                     degenerate_to_sgd=c.degenerate_to_sgd,
+                     rectify=c.rectify)
+
+
+@option(OptimizerConfigs.optimizer, 'Noam')
+def _noam_optimizer(c: OptimizerConfigs):
+    from labml_nn.optimizers.noam import Noam
+    return Noam(c.parameters,
+                lr=c.learning_rate, betas=c.betas, eps=c.eps,
+                weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad, warmup=c.warmup,
+                d_model=c.d_model)
+
+
+@option(OptimizerConfigs.optimizer, 'Sophia')
+def _sophia_optimizer(c: OptimizerConfigs):
+    from labml_nn.optimizers.sophia import Sophia
+    return Sophia(c.parameters,
+                  lr=c.learning_rate, betas=c.betas, eps=c.eps,
+                  weight_decay=c.weight_decay_obj, rho=c.rho)
+
+
+@option(OptimizerConfigs.optimizer, 'AdamWarmupCosineDecay')
+def _noam_optimizer(c: OptimizerConfigs):
+    from labml_nn.optimizers.adam_warmup_cosine_decay import AdamWarmupCosineDecay
+    return AdamWarmupCosineDecay(c.parameters,
+                                 lr=c.learning_rate, betas=c.betas, eps=c.eps,
+                                 weight_decay=c.weight_decay_obj, amsgrad=c.amsgrad,
+                                 warmup=c.warmup, total_steps=c.total_steps)