prpy 0.2.2__py3-none-any.whl

prpy/tensorflow/model_saver.py

@@ -0,0 +1,203 @@
+# Copyright (c) 2024 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""Save the best models"""
+
+import glob
+import logging
+import os
+import shutil
+import tensorflow as tf
+from typing import Union, Callable
+
+class Candidate(object):
+  """A candidate model with a score"""
+  def __init__(
+      self,
+      score: Union[float, int],
+      dir: str,
+      filename: str
+    ):
+    """Initialise a new candidate.
+
+    Args:
+      score: The score achieved by the candidate
+      dir: The directory where the candidate model can be saved
+      filename: The filename under which the candidate model can be saved
+    """
+    assert isinstance(score, (float, int))
+    assert isinstance(dir, str)
+    assert isinstance(filename, str)
+    self.score = score
+    self.filepath = os.path.join(dir, filename)
+
+class ModelSaver(object):
+  """Save the best models to disk"""
+  def __init__(
+      self,
+      dir: str = "checkpoints",
+      keep_best: int = 5,
+      keep_latest: int = 1,
+      save_format: str = "tf",
+      save_optimizer: bool = False,
+      compare_fn: Callable[[float, float], bool] = lambda x,y: x.score < y.score,
+      sort_reverse: bool = False,
+      log_fn: Callable[[str], None] = logging.info
+    ):
+    """Init the ModelSaver
+    
+    Args:
+      dir: The directory where models should be saved
+      keep_best: The number of best scoring models to keep
+      keep_latest: The number of latest models to keep
+      save_format: Model format for saving ['tf' or 'h5']
+      save_optimizer: Also save optimizer state?
+      compare_fn: Function that compares two scores
+      sort_reverse: Reverse sort order?
+      log_fn: Function to write logs
+    """
+    assert isinstance(dir, str)
+    assert isinstance(keep_best, int)
+    assert isinstance(keep_latest, int)
+    assert isinstance(save_format, str)
+    assert isinstance(save_optimizer, bool)
+    assert callable(compare_fn)
+    assert isinstance(sort_reverse, bool)
+    assert callable(log_fn)
+    self.best_candidates = []
+    self.latest_candidates = []
+    # The destination directory (make if necessary)
+    self.dir = dir
+    if not os.path.exists(self.dir):
+      os.makedirs(self.dir)
+    self.keep_best = keep_best
+    self.keep_latest = keep_latest
+    self.save_format = save_format
+    self.save_optimizer = save_optimizer
+    self.compare_fn_best = compare_fn
+    self.compare_fn_latest = lambda x,y: x.score > y.score
+    self.sort_reverse = sort_reverse
+    self.log_fn = log_fn
+
+  def __save(self, model: tf.keras.Model, filepath: str):
+    """Save a model to disk.
+    
+    Args:
+      model: The keras model to be saved
+      filepath: The filepath to save to
+    """
+    assert isinstance(model, tf.keras.Model)
+    assert isinstance(filepath, str)
+    if self.save_format == 'h5': filepath += '.h5'
+    # Save model
+    if self.save_optimizer:
+      model.save(
+        filepath=filepath, overwrite=True, include_optimizer=True,
+        save_format=self.save_format)
+    else:
+      model.save_weights(
+        filepath=filepath, overwrite=True, save_format=self.save_format)
+
+  def save_keep(self, model: tf.keras.Model, step: int, name: str):
+    """Save and keep the given model.
+
+    Args:
+      model: The model to be saved and kept
+      step: The current training step
+      name: The model name
+    """
+    assert isinstance(model, tf.keras.Model)
+    assert isinstance(step, int)
+    assert isinstance(name, str)
+    self.log_fn("Saving and keeping model for step {}".format(step))
+    filepath = os.path.join(self.dir, name + "_keep_" + str(step))
+    self.__save(model=model, filepath=filepath)
+
+  def save_latest(self, model: tf.keras.Model, step: int, name: str):
+    """Save the given model as currently latest.
+    
+    Args:
+      model: The model to be saved as latest
+      step: The current training step
+      name: The model name
+    """
+    assert isinstance(model, tf.keras.Model)
+    assert isinstance(step, int)
+    assert isinstance(name, str)
+    name = name + "_latest_" + str(step)
+    # Use step as score
+    candidate = Candidate(score=step, dir=self.dir, filename=name)
+    if len(self.latest_candidates) < self.keep_latest \
+      or self.compare_fn_latest(candidate, self.latest_candidates[-1]):
+      self.log_fn("Saving latest model for step {}".format(step))
+      # Keep candidate
+      self.latest_candidates.append(candidate)
+      self.latest_candidates = sorted(
+        self.latest_candidates, key=lambda x: x.score, reverse=True)
+      # Save candidate
+      self.__save(model, filepath=candidate.filepath)
+      # Prune candidate
+      for candidate in self.latest_candidates[self.keep_latest:]:
+        for file in glob.glob(r'{}*'.format(candidate.filepath)):
+          if self.save_format == 'tf' and self.save_optimizer:
+            shutil.rmtree(file)
+          else:
+            os.remove(file)
+      self.latest_candidates = self.latest_candidates[0:self.keep_latest]
+
+  def save_best(self, model: tf.keras.Model, score: float, step: int, name: str):
+    """Save the given model as a candidate for best model.
+
+    Args:
+      model: The model to be saved as candidate for best model
+      score: The score achieved by the model
+      step: The current training step
+      name: The name of the model
+    """
+    assert isinstance(model, tf.keras.Model)
+    assert isinstance(score, float)
+    assert isinstance(step, int)
+    assert isinstance(name, str)
+    self.log_fn('Saving model for step {}'.format(step))
+    name = name + "_best_" + str(step)
+    candidate = Candidate(score, dir=self.dir, filename=name)
+    if len(self.best_candidates) < self.keep_best \
+      or self.compare_fn_best(candidate, self.best_candidates[-1]):
+      # Keep candidate
+      self.log_fn("Keeping model {} with score {:.4f}".format(
+        candidate.filepath, candidate.score))
+      self.best_candidates.append(candidate)
+      self.best_candidates = sorted(
+        self.best_candidates, key=lambda x: x.score, reverse=self.sort_reverse)
+      # Save candidate
+      self.__save(model, filepath=candidate.filepath)
+      # Prune candidates
+      for candidate in self.best_candidates[self.keep_best:]:
+        self.log_fn('Removing old model {} with score {:.4f}'.format(
+          candidate.filepath, candidate.score))
+        for file in glob.glob(r'{}*'.format(candidate.filepath)):
+          if self.save_format == 'tf' and self.save_optimizer:
+            shutil.rmtree(file)
+          else:
+            os.remove(file)
+      self.best_candidates = self.best_candidates[0:self.keep_best]
+    else:
+      # Skip the candidate
+      self.log_fn('Skipping candidate {}'.format(candidate.filepath))

prpy/tensorflow/nan.py

@@ -0,0 +1,226 @@
+# Copyright (c) 2024 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import tensorflow as tf
+from typing import Union, Tuple, Callable
+
+def reduce_nanmean(
+    x: tf.Tensor,
+    axis: Union[int, tuple, None] = None
+  ) -> tf.Tensor:
+  """tf.reduce_mean, ignoring non-finite vals.
+
+  - Returns `nan` for all-nan slices.
+  
+  Args:
+    x: The input tensor.
+    axis: The dimension to reduce.
+  Returns:
+    The reduced tensor.
+  """
+  assert isinstance(x, tf.Tensor)
+  assert axis is None or isinstance(axis, (int, tuple))
+  mask = tf.math.is_finite(x)
+  numerator = tf.reduce_sum(tf.where(mask, x, tf.zeros_like(x)), axis=axis)
+  denominator = tf.reduce_sum(tf.cast(mask, dtype=x.dtype), axis=axis)
+  return numerator / denominator
+
+class ReduceNanMean:
+  """tf.reduce_mean, ignoring non-finite values. Supports gradient.
+
+  Behavior when x is non-finite:
+  - out: Non-finite vals in a slice contribute 0
+  - out: All-non-finite slices are nan
+  - grad = 0
+  """
+  def __init__(self, axis: Union[int, tuple, None] = None):
+    """Initialize.
+
+    Args:
+      axis: Axes to reduce by mean
+    """
+    assert axis is None or isinstance(axis, int) or (isinstance(axis, tuple) and all(isinstance(i, int) for i in axis))
+    self.axis = axis
+  @tf.custom_gradient
+  def __call__(self, x: tf.Tensor) -> Tuple[tf.Tensor, Callable[[tf.Tensor], tf.Tensor]]:
+    """Compute the mean.
+
+    Args:
+      x: The values.
+    Returns:
+      out: The computed mean.
+      grad: Function calculating the gradient
+    """
+    assert isinstance(x, tf.Tensor)
+    mask = tf.math.is_finite(x)
+    num = tf.reduce_sum(tf.where(mask, x, tf.zeros_like(x)), axis=self.axis)
+    den = tf.reduce_sum(tf.cast(mask, dtype=x.dtype), axis=self.axis)
+    mean = num / den
+    def grad(upstream: tf.Tensor) -> tf.Tensor:
+      den = tf.reduce_sum(tf.cast(mask, dtype=x.dtype), axis=self.axis)
+      # Tile upstream to match x
+      if self.axis is not None:
+        axis_list = list(self.axis) if isinstance(self.axis, tuple) else [self.axis]
+        # Expand dims to match x 
+        for axis in axis_list:
+          upstream = tf.expand_dims(upstream, axis=axis)
+          den = tf.expand_dims(den, axis=axis)
+        # Tile
+        tile_shape = [1 if s1 == s2 else s2 if s1 == 1 else s1 for s1, s2 in zip(x.shape, upstream.shape)]
+        upstream = tf.tile(upstream, tile_shape)
+        den = tf.tile(den, tile_shape)      
+      # Compute gradient and set to 0 where input was not finite
+      dout_dx = tf.where(mask, upstream / den, tf.zeros_like(x))
+      return dout_dx
+    return mean, grad
+
+def reduce_nansum(
+    x: tf.Tensor,
+    weight: Union[tf.Tensor, list, None] = None,
+    axis: Union[int, tuple, None] = None,
+    default: Union[float, int] = float('nan')
+  ) -> tf.Tensor:
+  """tf.reduce_sum, weighted by weight, ignoring non-finite values.
+
+  - Returns default for all-nan slices.
+  
+  Args:
+    x: The input tensor.
+    weight: The weight tensor, with the same shape as x or broadcastable to it.
+    axis: The dimension to reduce.
+    default: The value to return for all-non-finite slices.
+  Returns:
+    The reduced tensor.
+  """
+  assert isinstance(x, tf.Tensor)
+  assert weight is None or isinstance(weight, (tf.Tensor, list))
+  assert axis is None or isinstance(axis, (int, tuple))
+  assert isinstance(default, (float, int))
+  mask = tf.math.is_finite(x)
+  if weight is None:
+    sum = tf.reduce_sum(tf.where(mask, x, tf.zeros_like(x)), axis=axis)
+  else:
+    weight = tf.where(mask, weight, tf.zeros_like(weight))
+    sum = tf.reduce_sum(tf.where(mask, x * tf.cast(weight, x.dtype), tf.zeros_like(x)), axis=axis)
+  # If there are no finite elements in a slice, return the default.
+  return tf.where(tf.reduce_all(tf.logical_not(mask), axis=axis),
+                  tf.constant(default, dtype=x.dtype),
+                  sum)
+
+class ReduceNanSum:
+  """tf.reduce_sum, weighted by weight, ignoring non-finite values. Supports gradient.
+
+  Behavior when x is non-finite:
+  - out: Non-finite vals in a slice contribute 0
+  - out: All-non-finite slices are set to default value
+  - grad = 0
+  """
+  def __init__(
+      self,
+      weight: Union[tf.Tensor, list, None] = None,
+      axis: Union[int, tuple, None] = None,
+      default: Union[int, float] = float('nan')
+    ):
+    """Initialize.
+
+    Args:
+      weight: The weight tensor, with the same shape as x or broadcastable to it.
+      axis: Axes to reduce by sum
+      default: Value for all-non-finite slices
+    """
+    assert weight is None or isinstance(weight, (tf.Tensor, list))
+    assert axis is None or isinstance(axis, (int, tuple))
+    assert isinstance(default, (int, float))
+    self.weight = weight
+    self.axis = axis
+    self.default = default
+  @tf.custom_gradient
+  def __call__(self, x: tf.Tensor) -> Tuple[tf.Tensor, Callable[[tf.Tensor], tf.Tensor]]:
+    """Compute the sum.
+
+    Args:
+      x: The values.
+    Returns:
+      out: The computed sum.
+      grad: Function calculating the gradient
+    """
+    assert isinstance(x, tf.Tensor)
+    mask = tf.math.is_finite(x)
+    if self.weight is None:
+      sum = tf.reduce_sum(tf.where(mask, x, tf.zeros_like(x)), axis=self.axis)
+    else:
+      weight = tf.where(mask, self.weight, tf.zeros_like(self.weight))
+      sum = tf.reduce_sum(tf.where(mask, x * tf.cast(weight, x.dtype), tf.zeros_like(x)), axis=self.axis)
+    # If there are no finite elements in a slice, return the default.
+    out = tf.where(tf.reduce_all(tf.logical_not(mask), axis=self.axis),
+                   tf.cast(self.default, x.dtype),
+                   sum)
+    def grad(upstream: tf.Tensor) -> tf.Tensor:
+      # Tile upstream to match x
+      if self.axis is not None:
+        axis_list = list(self.axis) if isinstance(self.axis, tuple) else [self.axis]
+        # Expand dims to match x 
+        for axis in axis_list:
+          upstream = tf.expand_dims(upstream, axis=axis)
+        # Tile
+        tile_shape = [1 if s1 == s2 else s2 if s1 == 1 else s1 for s1, s2 in zip(x.shape, upstream.shape)]
+        upstream = tf.tile(upstream, tile_shape)
+      # Set the gradient to 0 where input was not finite
+      dout_dx = tf.where(mask, upstream, tf.zeros_like(x))
+      return dout_dx
+    return out, grad
+
+class NanLinearCombination:
+  """Linear combination with one fixed element. Supports gradient.
+
+  - Calculates (1 - x) * val_1 + x * val_2
+  - Behavior when val_1 or val_2 are non-finite:
+    - out = nan
+    - grad = 0
+  """
+  @tf.custom_gradient
+  def __call__(
+    self,
+    x: tf.Tensor,
+    val_1: tf.Tensor,
+    val_2: tf.Tensor
+  ) -> Tuple[tf.Tensor, Callable[[tf.Tensor], Tuple[tf.Tensor, tf.Tensor, tf.Tensor]]]:
+    """Compute the linear combination.
+
+    Args:
+      x: The combination weight. All elements must be in [0, 1]
+      val_1: The second value used in the linear combination.
+      val_2: The second value used in the linear combination.
+    Returns:
+      out: The computed sum.
+      grad: Function calculating the gradient
+    """
+    # Compute the linear combination
+    val_1 = tf.broadcast_to(val_1, x.shape)
+    val_2 = tf.broadcast_to(val_2, x.shape)
+    out = (1 - x) * val_1 + x * val_2
+    def grad(upstream: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
+      mask = tf.math.logical_and(tf.math.is_finite(val_1), tf.math.is_finite(val_2))
+      dout_dx = tf.where(mask, upstream * (val_2 - val_1), tf.zeros_like(x))
+      dout_dval_1 = upstream * (1 - x)
+      dout_dval_2 = upstream * x
+      return dout_dx, dout_dval_1, dout_dval_2
+    return out, grad
+  

prpy/tensorflow/optimizer.py

@@ -0,0 +1,233 @@
+# Copyright (c) 2024 Philipp Rouast
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from packaging import version
+import tensorflow as tf
+
+if version.parse(tf.__version__) <= version.parse("2.6.5"):
+  from tensorflow.keras.optimizers import Adam # Legacy / V2 optimizer
+  from tensorflow.keras.mixed_precision import LossScaleOptimizer
+elif version.parse(tf.__version__) >= version.parse("2.9") and version.parse(tf.__version__) <= version.parse("2.11"):
+  from keras.optimizers.optimizer_experimental.adam import Adam # Experimental / V3 optimizer
+  from keras.optimizers.optimizer_experimental.adamw import AdamW # Experimental / V3 optimizer
+  from keras.mixed_precision.loss_scale_optimizer import LossScaleOptimizerV3 as LossScaleOptimizer
+elif version.parse(tf.__version__) > version.parse("2.11"):
+  from tensorflow.keras.optimizers import Adam, AdamW # V3 optimizer
+  from keras.mixed_precision.loss_scale_optimizer import LossScaleOptimizerV3 as LossScaleOptimizer
+else:
+  raise ImportError("This version of TensorFlow is not compatible.")
+
+class EpochAdamMetaclass(type):
+  """Metaclass that delegates EpochAdam instance creation."""
+  def __call__(cls, **kwargs):
+    if version.parse(tf.__version__) <= version.parse("2.6.5"):
+      return EpochAdamV2(**kwargs)
+    else:
+      return EpochAdamExperimental(**kwargs)
+
+class EpochAdamWMetaclass(type):
+  """Metaclass that delegates EpochAdamW instance creation."""
+  def __call__(cls, **kwargs):
+    if version.parse(tf.__version__) > version.parse("2.9"):
+      return EpochAdamWExperimental(**kwargs)
+    else:
+      raise ImportError("This version of TensorFlow is not compatible.")
+
+class EpochAdam(metaclass=EpochAdamMetaclass):
+  pass
+
+class EpochAdamW(metaclass=EpochAdamWMetaclass):
+  pass
+
+class EpochAdamExperimental(Adam):
+  """Experimental Adam optimizer that retrieves learning rate based on epochs"""
+  def __init__(self, **kwargs):
+    # Create epochs counter variable
+    with tf.init_scope():
+      # Lift the variable creation to init scope to avoid environment issue.
+      self._epochs = tf.Variable(
+        0, name="epochs", dtype=tf.int64, trainable=False)
+    super().__init__(**kwargs)
+    self._variables.append(self._epochs)
+  @property
+  def epochs(self):
+    return self._epochs
+  @epochs.setter
+  def epochs(self, variable):
+    if getattr(self, "_built", False):
+      raise RuntimeError(
+          "Cannot set `epochs` to a new Variable after "
+          "the Optimizer weights have been created. Here it is "
+          f"attempting to set `iterations` to {variable}."
+          "Usually this means you are trying to set `iterations`"
+          " after calling `apply_gradients()`. Please set "
+          "`iterations` before calling `apply_gradients()`.")
+    self._epochs = variable 
+  def _build_learning_rate(self, learning_rate):
+    with tf.init_scope():
+      if isinstance(learning_rate, tf.keras.optimizers.schedules.LearningRateSchedule):
+        # Create a variable to hold the current learning rate.
+        current_learning_rate = tf.convert_to_tensor(
+            learning_rate(self.epochs))
+        self._current_learning_rate = tf.Variable(
+            current_learning_rate,
+            name="current_learning_rate",
+            dtype=current_learning_rate.dtype,
+            trainable=False)
+        return learning_rate
+      return tf.Variable(
+        learning_rate,
+        name="learning_rate",
+        dtype=tf.float32,
+        trainable=False)
+  def _compute_current_learning_rate(self):
+    if isinstance(self._learning_rate, tf.keras.optimizers.schedules.LearningRateSchedule):
+      # Compute the current learning rate at the beginning of variable update.
+      if hasattr(self, "_current_learning_rate"):
+        self._current_learning_rate.assign(
+            self._learning_rate(self.epochs))
+      else:
+        current_learning_rate = tf.convert_to_tensor(
+          self._learning_rate(self.epochs))
+        self._current_learning_rate = tf.Variable(
+          current_learning_rate,
+          name="current_learning_rate",
+          dtype=current_learning_rate.dtype,
+          trainable=False)
+  def finish_epoch(self):
+    """Increment epoch count and re-compute lr"""
+    self._epochs.assign_add(1)
+    self._compute_current_learning_rate()
+
+class EpochAdamWExperimental(AdamW):
+  """Experimental AdamW optimizer that retrieves learning rate based on epochs"""
+  def __init__(self, **kwargs):
+    # Create epochs counter variable
+    with tf.init_scope():
+      # Lift the variable creation to init scope to avoid environment issue.
+      self._epochs = tf.Variable(
+        0, name="epochs", dtype=tf.int64, trainable=False)
+    super().__init__(**kwargs)
+    self._variables.append(self._epochs)
+  @property
+  def epochs(self):
+    return self._epochs
+  @epochs.setter
+  def epochs(self, variable):
+    if getattr(self, "_built", False):
+      raise RuntimeError(
+          "Cannot set `epochs` to a new Variable after "
+          "the Optimizer weights have been created. Here it is "
+          f"attempting to set `iterations` to {variable}."
+          "Usually this means you are trying to set `iterations`"
+          " after calling `apply_gradients()`. Please set "
+          "`iterations` before calling `apply_gradients()`.")
+    self._epochs = variable 
+  def _build_learning_rate(self, learning_rate):
+    with tf.init_scope():
+      if isinstance(learning_rate, tf.keras.optimizers.schedules.LearningRateSchedule):
+        # Create a variable to hold the current learning rate.
+        current_learning_rate = tf.convert_to_tensor(
+            learning_rate(self.epochs))
+        self._current_learning_rate = tf.Variable(
+            current_learning_rate,
+            name="current_learning_rate",
+            dtype=current_learning_rate.dtype,
+            trainable=False)
+        return learning_rate
+      return tf.Variable(
+        learning_rate,
+        name="learning_rate",
+        dtype=tf.float32,
+        trainable=False)
+  def _compute_current_learning_rate(self):
+    if isinstance(self._learning_rate, tf.keras.optimizers.schedules.LearningRateSchedule):
+      # Compute the current learning rate at the beginning of variable update.
+      if hasattr(self, "_current_learning_rate"):
+        self._current_learning_rate.assign(
+            self._learning_rate(self.epochs))
+      else:
+        current_learning_rate = tf.convert_to_tensor(
+          self._learning_rate(self.epochs))
+        self._current_learning_rate = tf.Variable(
+          current_learning_rate,
+          name="current_learning_rate",
+          dtype=current_learning_rate.dtype,
+          trainable=False)
+  def finish_epoch(self):
+    """Increment epoch count and re-compute lr"""
+    self._epochs.assign_add(1)
+    self._compute_current_learning_rate()
+
+class EpochAdamV2(Adam):
+  """V2 Adam optimizer that retrieves learning rate based on epochs"""
+  def __init__(self, **kwargs):
+    super().__init__(**kwargs)
+    self._epochs = None
+  def _decayed_lr(self, var_dtype):
+    """Get learning rate based on epochs."""
+    lr_t = self._get_hyper("learning_rate", var_dtype)
+    if isinstance(lr_t, tf.keras.optimizers.schedules.LearningRateSchedule):
+      epochs = tf.cast(self.epochs, var_dtype)
+      lr_t = tf.cast(lr_t(epochs), var_dtype)
+    return lr_t
+  @property
+  def epochs(self):
+    """Variable. The number of epochs."""
+    if self._epochs is None:
+      self._epochs = self.add_weight(
+        "epochs", shape=[], dtype=tf.int64, trainable=False,
+        aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA)
+      self._weights.append(self._epochs)
+    return self._epochs
+  def finish_epoch(self):
+    """Increment epoch count"""
+    return self._epochs.assign_add(1)
+
+class EpochLossScaleOptimizerMetaclass(type):
+  """Metaclass that delegates EpochLossScaleOptimizer instance creation."""
+  def __call__(cls, inner_optimizer, **kwargs):
+    if version.parse(tf.__version__) <= version.parse("2.6.5"):
+      return EpochLossScaleOptimizerV2(inner_optimizer, **kwargs)
+    else:
+      return EpochLossScaleOptimizerV3(inner_optimizer, **kwargs)
+
+class EpochLossScaleOptimizer(metaclass=EpochLossScaleOptimizerMetaclass):
+  pass
+
+class EpochLossScaleOptimizerV2(LossScaleOptimizer):
+  """Subclass LossScaleOptimizer to use epochs. Works for tensorflow<=2.6.5"""
+  @property
+  def epochs(self):
+    """Variable. The number of epochs."""
+    return self._optimizer.epochs
+  def finish_epoch(self):
+    """Increment epoch count"""
+    return self._optimizer.finish_epoch()
+
+class EpochLossScaleOptimizerV3(LossScaleOptimizer):
+  """Subclass LossScaleOptimizer to use epochs."""
+  @property
+  def epochs(self):
+    """Variable. The number of epochs."""
+    return self._optimizer.epochs
+  def finish_epoch(self):
+    """Increment epoch count"""
+    return self._optimizer.finish_epoch()