Construct a new LAMB optimizer. Args: learning_rate: A `Tensor` or a floating point value. or a schedule that is a `tf_keras.optimizers.schedules.LearningRateSchedule` The learning rate. beta_1: A `float` value or a constant `float` tensor. The exponential
(
self,
learning_rate: Union[FloatTensorLike, Callable] = 0.001,
beta_1: FloatTensorLike = 0.9,
beta_2: FloatTensorLike = 0.999,
epsilon: FloatTensorLike = 1e-6,
weight_decay_rate: FloatTensorLike = 0.0,
exclude_from_weight_decay: Optional[List[str]] = None,
exclude_from_layer_adaptation: Optional[List[str]] = None,
name: str = "LAMB",
**kwargs,
)
| 35 | """ |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | learning_rate: Union[FloatTensorLike, Callable] = 0.001, |
| 40 | beta_1: FloatTensorLike = 0.9, |
| 41 | beta_2: FloatTensorLike = 0.999, |
| 42 | epsilon: FloatTensorLike = 1e-6, |
| 43 | weight_decay_rate: FloatTensorLike = 0.0, |
| 44 | exclude_from_weight_decay: Optional[List[str]] = None, |
| 45 | exclude_from_layer_adaptation: Optional[List[str]] = None, |
| 46 | name: str = "LAMB", |
| 47 | **kwargs, |
| 48 | ): |
| 49 | """Construct a new LAMB optimizer. |
| 50 | |
| 51 | Args: |
| 52 | learning_rate: A `Tensor` or a floating point value. or a schedule that |
| 53 | is a `tf_keras.optimizers.schedules.LearningRateSchedule` The learning |
| 54 | rate. |
| 55 | beta_1: A `float` value or a constant `float` tensor. The exponential |
| 56 | decay rate for the 1st moment estimates. |
| 57 | beta_2: A `float` value or a constant `float` tensor. The exponential |
| 58 | decay rate for the 2nd moment estimates. |
| 59 | epsilon: A small constant for numerical stability. |
| 60 | weight_decay_rate: weight decay rate. |
| 61 | exclude_from_weight_decay: List of regex patterns of variables excluded |
| 62 | from weight decay. Variables whose name contain a substring matching |
| 63 | the pattern will be excluded. |
| 64 | exclude_from_layer_adaptation: List of regex patterns of variables |
| 65 | excluded from layer adaptation. Variables whose name contain a |
| 66 | substring matching the pattern will be excluded. |
| 67 | name: Optional name for the operations created when applying gradients. |
| 68 | Defaults to "LAMB". |
| 69 | **kwargs: keyword arguments. Allowed to be {`clipnorm`, `clipvalue`, |
| 70 | `lr`, `decay`}. `clipnorm` is clip gradients by norm; `clipvalue` is |
| 71 | clip gradients by value, `decay` is included for backward |
| 72 | compatibility to allow time inverse decay of learning rate. `lr` is |
| 73 | included for backward compatibility, recommended to use |
| 74 | `learning_rate` instead. |
| 75 | """ |
| 76 | super().__init__(name, **kwargs) |
| 77 | |
| 78 | # Just adding the square of the weights to the loss function is *not* |
| 79 | # the correct way of using L2 regularization/weight decay with Adam, |
| 80 | # since that will interact with the m and v parameters in strange ways. |
| 81 | # |
| 82 | # Instead we want to decay the weights in a manner that doesn't interact |
| 83 | # with the m/v parameters. |
| 84 | self._set_hyper("weight_decay_rate", weight_decay_rate) |
| 85 | self._set_hyper("learning_rate", kwargs.get("lr", learning_rate)) |
| 86 | |
| 87 | # This is learning rate decay for using keras learning rate schedule. |
| 88 | self._set_hyper("decay", self._initial_decay) |
| 89 | self._set_hyper("beta_1", beta_1) |
| 90 | self._set_hyper("beta_2", beta_2) |
| 91 | self.epsilon = epsilon or tf.backend_config.epsilon() |
| 92 | self.exclude_from_weight_decay = exclude_from_weight_decay |
| 93 | # exclude_from_layer_adaptation is set to exclude_from_weight_decay if |
| 94 | # the arg is None. |