Internal change

PiperOrigin-RevId: 264474346

official/bert/optimization.py

@@ -105,12 +105,14 @@ class AdamWeightDecay(tf.keras.optimizers.Adam):
                epsilon=1e-7,
                amsgrad=False,
                weight_decay_rate=0.0,
+               include_in_weight_decay=None,
                exclude_from_weight_decay=None,
                name='AdamWeightDecay',
                **kwargs):
     super(AdamWeightDecay, self).__init__(
         learning_rate, beta_1, beta_2, epsilon, amsgrad, name, **kwargs)
     self.weight_decay_rate = weight_decay_rate
+    self._include_in_weight_decay = include_in_weight_decay
     self._exclude_from_weight_decay = exclude_from_weight_decay
 
   @classmethod
@@ -178,6 +180,12 @@ class AdamWeightDecay(tf.keras.optimizers.Adam):
     """Whether to use L2 weight decay for `param_name`."""
     if self.weight_decay_rate == 0:
       return False
+
+    if self._include_in_weight_decay:
+      for r in self._include_in_weight_decay:
+        if re.search(r, param_name) is not None:
+          return True
+
     if self._exclude_from_weight_decay:
       for r in self._exclude_from_weight_decay:
         if re.search(r, param_name) is not None: