prpy 0.2.2__py3-none-any.whl

prpy/numpy/stride_tricks.py

@@ -0,0 +1,182 @@
+# 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 numpy as np
+from typing import Union
+
+def window_view(
+    x: np.ndarray,
+    min_window_size: int,
+    max_window_size: int,
+    overlap: int,
+    pad_mode: str = 'constant',
+    const_val: Union[float, int] = np.nan
+  ) -> np.ndarray:
+  """Create a window view into an n-d array `x` along its first dim.
+
+  Args:
+    x: The n-d array into which we want to create a windowed view
+    min_window_size: The minimum window size
+    max_window_size: The maximum window size
+    overlap: The overlap of the sliding windows
+    pad_mode: The pad mode
+    const_val: The constant value to be padded with if pad_mode == 'constant'
+  Returns:
+    y: The n+1-d windowed view into x of shape (n_windows, window_size, ...)
+    pad_start: How much padding was applied at the start (scalar)
+    pad_end: How much padding was applied at the end (scalar)
+  """
+  assert isinstance(min_window_size, int) and min_window_size >= 0
+  assert isinstance(max_window_size, int) and max_window_size >= min_window_size
+  assert isinstance(overlap, int) and overlap < max_window_size, "overlap must be smaller than max_window_size"
+  assert isinstance(pad_mode, str)
+  assert isinstance(const_val, (float, int))
+  x = np.asarray(x)
+  original_len = x.shape[0]
+  step_size = max_window_size - overlap
+  # Pad the start using variable window sizes
+  pad_start = max_window_size - min_window_size
+  # Pad the end if there is a remainder
+  remainder = (original_len - min_window_size) % step_size
+  pad_end = 0 if remainder == 0 else step_size - remainder
+  pad_width = ((pad_start, pad_end),) + ((0, 0),) * (x.ndim - 1)
+  if pad_mode=='constant':
+    x = np.pad(x, pad_width, mode=pad_mode, constant_values=(const_val,))
+  else:
+    x = np.pad(x, pad_width, mode=pad_mode)
+  # Calculate the shape of the view for the first dimension
+  new_shape = ((pad_start + original_len + pad_end - max_window_size) // step_size + 1, max_window_size)
+  # Add shapes for following dimensions
+  new_shape += x.shape[1:]
+  # Calculate the stride of the view for the first dimension
+  new_strides = ((step_size * x.strides[0],) + (x.strides[0],))
+  # Add strides for following dimensions
+  new_strides += x.strides[1:]
+  # Generate the view
+  y = np.lib.stride_tricks.as_strided(
+    x, shape=new_shape, strides=new_strides)
+  # Return
+  return y, pad_start, pad_end
+
+def reduce_window_view(
+    x: np.ndarray,
+    overlap: int,
+    pad_end: int = 0,
+    hanning: bool = False
+  ) -> np.ndarray:
+  """Reduce an n-d window view by arranging the first dimension as sliding
+  windows and then reducing it using the mean.
+
+  Args:
+    x: The n-d window view of shape (n_windows, window_size, ...)
+    overlap: The overlap with which the window view was created
+    pad_end: How much padding was applied to the end when the window view was created
+    hanning: Whether to reduce the window view with hanning windows
+  Returns:
+    mean: The n-1-d reduced window view [original_len, ...]
+  """
+  assert isinstance(x, np.ndarray)
+  assert isinstance(hanning, bool)
+  # Infer number of windows and original length (including extra padding)
+  num_windows = x.shape[0]
+  window_size = x.shape[1]
+  assert isinstance(overlap, int) and overlap >= 0 and overlap < window_size
+  assert isinstance(pad_end, int) and pad_end >= 0 and pad_end < window_size
+  original_len_with_pad_end = num_windows * window_size - (num_windows-1) * overlap
+  # Apply hanning window to taper x
+  if hanning:
+    x *= np.hanning(window_size)
+  # Add padding to extend the matrix to the dimensions of the diagonal matrix
+  padding = ((0, 0), (0, original_len_with_pad_end - window_size)) + ((0, 0),) * (x.ndim - 2)
+  y = np.pad(x, padding, 'constant')
+  # Use as_strided to create a view of the diagonal matrix
+  #  that aligns the windows temporally as they were generated
+  # https://stackoverflow.com/a/60460462/3595278
+  y_roll = y[:, [*range(y.shape[1]),*range(y.shape[1]-1)]].copy() #need `copy`
+  view_strides = list(y_roll.strides)
+  view_strides.insert(1, y_roll.strides[1])
+  view_strides = tuple(view_strides)
+  view_shape = list(y.shape)
+  view_shape.insert(1, original_len_with_pad_end)
+  view_shape = tuple(view_shape)
+  view = np.lib.stride_tricks.as_strided(y_roll, view_shape, view_strides)
+  step_size = window_size - overlap
+  m = np.asarray([step_size * i for i in range(num_windows)])
+  view = view[np.arange(y.shape[0]), (original_len_with_pad_end-m)%original_len_with_pad_end]
+  # Merge the windows by taking the mean across the time dimension
+  mean = np.true_divide(
+    view.sum(0), np.maximum((view != 0).sum(0), 1))
+  # Trim result by pad_end if necessary
+  if pad_end > 0: mean = mean[:-pad_end]
+  return mean
+
+def resolve_1d_window_view(
+    x: np.ndarray,
+    window_size: int,
+    overlap: int,
+    pad_end: int,
+    fill_method: str
+  ) -> np.ndarray:
+  """Resolve an 1-d window view by extending it to the expected shape.
+  
+  - This is useful if processing on each window created a scalar value.
+
+  Args:
+    x: The 1-d window view to be resolved
+    window_size: The window size used to create the view
+    overlap: The overlap used to create the view
+    pad_end: How much padding was applied to the end
+    fill_method: Method for filling/padding the result
+  Returns:
+    vals: The 1-d resolved data
+  """
+  assert isinstance(x, np.ndarray) and len(x.shape) == 1
+  assert isinstance(window_size, int) and window_size > 0
+  assert isinstance(overlap, int) and overlap >= 0 and overlap < window_size
+  assert isinstance(pad_end, int) and pad_end >= 0 and pad_end < window_size
+  assert isinstance(fill_method, str)
+  if overlap == 0:
+    # If overlap is zero, we simply need to repeat each value to match window_size
+    vals = np.repeat(x, window_size)
+    if pad_end > 0:
+      # Trim end if it has been padded
+      vals = vals[:-pad_end]
+  elif overlap == window_size - 1:
+    # If overlap is one less than the window size, then values are mostly
+    # already ok. We just have to fill the start.
+    if fill_method == 'zero':
+      fill = 0.0
+    elif fill_method == 'mean':
+      fill = np.mean(x)
+    elif fill_method == 'start':
+      fill = x[0]
+    else:
+      raise ValueError("fill_method {} not supported".format(fill_method))
+    vals = np.concatenate([np.repeat(fill, window_size-1), x])
+  elif overlap < window_size - 1:
+    # For any other overlaps, we will have to build an intermediate 2-d
+    # window view of the vals and reduce it via reduce_window_view.
+    n = len(x)
+    x = np.reshape(np.repeat(x, window_size), (n, window_size))
+    vals = reduce_window_view(x, overlap=overlap)
+    if pad_end > 0:
+      # Trim end if it has been padded
+      vals = vals[:-pad_end]
+  return vals

prpy/tensorflow/__init__.py

@@ -0,0 +1,19 @@
+# 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.

prpy/tensorflow/image.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.
+
+import numpy as np
+import tensorflow as tf
+from typing import Union
+
+def _reduction_dims(
+    x: tf.Tensor,
+    axis: Union[int, tuple, None]
+  ) -> Union[int, tuple, tf.Tensor]:
+  """Resolve the reduction dims.
+
+  - If axis is None, returns all dims of x for reduction
+  
+  Args:
+    x: The tensor to be reduced
+    axis: Dims to be reduced or None
+  """
+  assert isinstance(x, tf.Tensor)
+  assert axis is None or isinstance(axis, int) or isinstance(axis, tuple)
+  if axis is not None:
+    return axis
+  else:
+    x_rank = None
+    if isinstance(x, tf.Tensor):
+      x_rank = x.shape.rank
+    # Fast path: avoid creating Rank and Range ops if ndims is known.
+    if x_rank:
+      return tf.constant(np.arange(x_rank, dtype=np.int32))
+    else:
+      # Otherwise, we rely on Range and Rank to do the right thing at run-time.
+      return tf.range(0, tf.rank(x))
+
+def standardize_images(
+    images: Union[tf.Tensor, np.ndarray],
+    axis: Union[int, tuple, None] = None
+  ) -> tf.Tensor:
+  """Standardize image data to have zero mean and unit variance.
+
+  Args:
+    images: The image data. 
+    axis: The dimensions to standardize across. Exclude the axes that should be
+      treated separately, e.g. the channel axis to treat each channel separately.
+      If None, standardize across all dimensions.
+  Returns:
+    images: The standardized image data as tf.float32.
+  """
+  assert isinstance(images, (tf.Tensor, np.ndarray))
+  assert axis is None or isinstance(axis, int) or isinstance(axis, tuple)
+  # Convert to tf.Tensor if necessary
+  if not tf.is_tensor(images):
+    images = tf.convert_to_tensor(images)
+  # We are working with tf.float32
+  images = tf.cast(images, dtype=tf.float32)
+  # Resolve axis arg
+  axis = _reduction_dims(images, axis)
+  # Compute the mean and std
+  num_pixels = tf.math.reduce_prod(tf.gather(tf.shape(images), axis))
+  mean = tf.math.reduce_mean(images, axis, keepdims=True)
+  std = tf.math.reduce_std(images, axis, keepdims=True)
+  # Apply a minimum normalization that protects us against uniform images
+  min_std = tf.math.rsqrt(tf.cast(num_pixels, dtype=tf.float32))
+  # Perform standardization
+  images = tf.subtract(images, mean)
+  images = tf.divide(images, tf.maximum(std, min_std))
+  return images
+
+def normalize_images(
+    images: Union[tf.Tensor, np.ndarray],
+    axis: Union[int, tuple, None] = None
+  ) -> tf.Tensor:
+  """Normalize image data to have zero mean.
+
+  Args:
+    images: The image data.
+    axis: The dimensions to normalize across. Exclude the axes that should be
+      treated separately, e.g. the channel axis to treat each channel separately.
+      If None, normalize across all dimensions.
+  Returns:
+    images: The normalized image data.
+  """
+  assert isinstance(images, (tf.Tensor, np.ndarray))
+  assert axis is None or isinstance(axis, int) or isinstance(axis, tuple)
+  # Convert to tf.Tensor if necessary
+  if not tf.is_tensor(images):
+    images = tf.convert_to_tensor(images)
+  # We are working with tf.float32
+  images = tf.cast(images, dtype=tf.float32)
+  # Resolve axis arg
+  axis = _reduction_dims(images, axis)
+  # Compute the mean
+  mean = tf.math.reduce_mean(images, axis, keepdims=True)
+  # Perform normalization
+  images = tf.subtract(images, mean)
+  return images
+
+def normalized_image_diff(
+    images: Union[tf.Tensor, np.ndarray],
+    axis: int = 0
+  ) -> tf.Tensor:
+  """Compute the normalized difference of adjacent images.
+
+  Args:
+    images: The image data as float32 in range [0, 1]
+    axis: Scalar, the dimension across which to calculate difference
+      normalization (e.g., the temporal/sequence dimension).
+  Returns:
+    images: The processed image data.
+  """
+  assert isinstance(images, (tf.Tensor, np.ndarray))
+  assert axis==0 or axis==1, "Only axis=0 or axis=1 supported"
+  # Convert to tf.Tensor if necessary
+  if not tf.is_tensor(images):
+    images = tf.convert_to_tensor(images)
+  # We are working with tf.float32
+  images = tf.cast(images, dtype=tf.float32)
+  diff = tf.cond(tf.equal(axis, 0),
+    true_fn=lambda: images[1:] - images[:-1],
+    false_fn=lambda: images[:,1:] - images[:,:-1])
+  sum = tf.cond(tf.equal(axis, 0),
+    true_fn=lambda: images[1:] + images[:-1],
+    false_fn=lambda: images[:,1:] + images[:,:-1])
+  sum = tf.clip_by_value(sum, clip_value_min=1e-7, clip_value_max=2)
+  return diff / sum
+
+def display_scale(
+    image: tf.Tensor
+  ) -> tf.Tensor:
+  """Scale any float32 image to the [0, 1] range for display.
+  
+  Args:
+    image: The image tensor.
+  Returns:
+    out: The scaled image for display.
+  """
+  assert isinstance(image, tf.Tensor)
+  min = tf.math.reduce_min(image)
+  range = tf.math.reduce_max(image) - min
+  return tf.clip_by_value((image-min) * 1.0/range, 0, 1)
+
+def resize_with_random_method(
+    images: tf.Tensor,
+    target_shape: Union[tuple, None] = (640, 640)
+  ) -> tf.Tensor:
+  """Resize image(s) with a random method.
+  
+  Args:
+    images: Image data, either shape (b, h, w, c) or (h, w, c)
+    target_shape: The resize shape in form (new_h, new_w)
+  Returns:
+    out: The resized image data with shape (b, new_h, new_w, c) or (new_h, new_w, c)
+  """
+  assert isinstance(images, tf.Tensor)
+  assert target_shape is None or (isinstance(target_shape, tuple) and len(target_shape) == 2 and all(isinstance(i, int) for i in target_shape))
+  # Draw a random number to determine resize method
+  resize_method = tf.random.uniform([], 0, 5, dtype=tf.int32)
+  def resize(method):
+    def _resize():
+      return tf.image.resize(
+        images, target_shape, method=method, antialias=True, preserve_aspect_ratio=False)
+    return _resize
+  # Resize using a random method
+  images = tf.case([(tf.equal(resize_method, 0), resize('bicubic')),
+                   (tf.equal(resize_method, 1), resize('area')),
+                   (tf.equal(resize_method, 2), resize('nearest')),
+                   (tf.equal(resize_method, 3), resize('lanczos3'))],
+                  default=resize('bilinear'))
+  return images
+
+def random_distortion(
+    images: tf.Tensor
+  ) -> tf.Tensor:
+  """Apply random distortion to image(s)
+
+  Args:
+    images: The image data (b, h, w, c) or (h, w, c)
+  Returns:
+    images: The distorted image data (b, h, w, c) or (h, w, c)
+  """
+  assert isinstance(images, tf.Tensor)
+  images = tf.image.random_brightness(images, 0.4)
+  images = tf.image.random_contrast(images, 0.5, 1.5)
+  images = tf.image.random_saturation(images, 0.5, 1.5)
+  images = tf.image.random_hue(images, 0.1)
+  return images

prpy/tensorflow/loss.py

@@ -0,0 +1,104 @@
+# 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
+
+def balanced_sample_weights(
+    labels: tf.Tensor,
+    unique: tf.Tensor
+  ) -> tf.Tensor:
+  """Calculate weights for a batch of dense categorical labels intended to be
+  multiplied with the losses
+  
+  - Larger weights for examples of under-represented classes, and smaller weights
+    for overrepresented classes, while keeping the total loss constant.
+  - Important: Only works for dense label representation that equal range from 0 to n
+  
+  Args:
+    labels: The dense categorical labels of shape (batch_size,) or (batch_size, 1)
+    unique: The unique labels of shape (n_unique_labels,)
+  Returns:
+    weights: The weights with same shape as labels.
+  """
+  assert isinstance(labels, tf.Tensor)
+  assert isinstance(unique, tf.Tensor)
+  # Remove empty dim if necessary
+  f_labels = tf.cast(tf.squeeze(labels), tf.int32)
+  # Determine count of labels in batch
+  # Cannot use bincount to be compatible with XLA
+  # count = tf.math.bincount(f_labels, minlength=n_labels)
+  def count(x):
+    return tf.reduce_sum(tf.cast(tf.equal(x, f_labels), tf.int32))
+  count = tf.map_fn(fn=count, elems=unique, fn_output_signature=tf.int32)
+  # Batch size and number of unique labels
+  batch_size = tf.size(f_labels)
+  unique_count = tf.reduce_sum(tf.cast(tf.math.greater(count, 0), tf.int32))
+  # Calculate the weight for each class
+  class_weights = tf.math.divide_no_nan(tf.cast(batch_size, tf.float32), tf.cast(count, tf.float32))
+  class_weights = class_weights / tf.cast(unique_count, tf.float32)
+  # Gather weights according to the actual categories of the batch elements
+  sample_weights = tf.gather(class_weights, f_labels)
+  # Reshape to original shape
+  sample_weights = tf.reshape(sample_weights, tf.shape(labels))
+  return sample_weights
+
+def smooth_l1_loss(
+    y_true: tf.Tensor,
+    y_pred: tf.Tensor,
+    keepdims: bool = False
+  ) -> tf.Tensor:
+  """Smooth L1 loss
+
+  Args:
+    y_true: Labels. Shape arbitrary.
+    y_pred: Predictions. Shape same as y_true.
+    keepdims: Keep original shape? Otherwise return global mean.
+  Returns:
+    loss: The loss. Shape same as original if keepdims, otherwise ()
+  """
+  assert isinstance(y_true, tf.Tensor)
+  assert isinstance(y_pred, tf.Tensor)
+  assert isinstance(keepdims, bool)
+  t = tf.abs(y_pred - y_true)
+  loss = tf.where(t < 1, 0.5 * t ** 2, t - 0.5)
+  return tf.cond(tf.equal(keepdims, tf.constant(True)),
+    true_fn=lambda: loss,
+    false_fn=lambda: tf.reduce_mean(loss))
+
+def mae_loss(
+    y_true: tf.Tensor,
+    y_pred: tf.Tensor,
+    keepdims: bool = False
+  ) -> tf.Tensor:
+  """Mean absolute error loss
+
+  Args:
+    y_true: Labels. Shape arbitrary.
+    y_pred: Predictions. Shape same as y_true.
+    keepdims: Keep batch dim? Otherwise return global mean.
+  Returns:
+    loss: The loss. Shape: If keepdims, original shape without final dim, otherwise ()
+  """
+  assert isinstance(y_true, tf.Tensor)
+  assert isinstance(y_pred, tf.Tensor)
+  assert isinstance(keepdims, bool)
+  return tf.cond(tf.equal(keepdims, tf.constant(True)),
+    true_fn=lambda: tf.reduce_mean(tf.abs(y_pred - y_true), axis=-1),
+    false_fn=lambda: tf.reduce_mean(tf.abs(y_pred - y_true)))

prpy/tensorflow/lr_schedule.py

@@ -0,0 +1,74 @@
+# 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
+
+class PiecewiseConstantDecayWithWarmup(tf.keras.optimizers.schedules.LearningRateSchedule):
+  """Piecewise constant decay with warmup."""
+  def __init__(
+      self,
+      boundaries: list,
+      values: list,
+      warmup_init_lr: Union[float, int],
+      warmup_steps: Union[int, tf.Variable],
+      name: Union[str, None] = None
+    ):
+    super(PiecewiseConstantDecayWithWarmup, self).__init__()
+    assert isinstance(boundaries, list)
+    assert isinstance(values, list)
+    assert isinstance(warmup_init_lr, (float, int))
+    assert isinstance(warmup_steps, (int, tf.Variable))
+    if len(boundaries) != len(values) - 1:
+        raise ValueError("The length of boundaries should be 1 less than the length of values")
+    self.boundaries = boundaries
+    self.values = values
+    self.name = name
+    self.warmup_steps = warmup_steps
+    self.warmup_init_lr = warmup_init_lr
+  def __call__(
+      self,
+      step: int
+    ) -> float:
+    assert isinstance(step, (int, tf.Variable))
+    with tf.name_scope(self.name or "PiecewiseConstantWarmUp"):
+      step = tf.cast(tf.convert_to_tensor(step), tf.float32)
+      pred_fn_pairs = []
+      warmup_steps = self.warmup_steps
+      boundaries = self.boundaries
+      values = self.values
+      warmup_init_lr = self.warmup_init_lr
+      pred_fn_pairs.append((step <= warmup_steps, lambda: warmup_init_lr + step * (values[0] - warmup_init_lr) / warmup_steps))
+      pred_fn_pairs.append((tf.logical_and(step <= boundaries[0], step > warmup_steps), lambda: tf.constant(values[0])))
+      pred_fn_pairs.append((step > boundaries[-1], lambda: tf.constant(values[-1])))
+      for low, high, v in zip(boundaries[:-1], boundaries[1:], values[1:-1]):
+        pred = (step > low) & (step <= high)
+        pred_fn_pairs.append((pred, lambda v=v: tf.constant(v)))
+      # The default isn't needed here because our conditions are mutually
+      # exclusive and exhaustive, but tf.case requires it.
+      return tf.case(pred_fn_pairs, lambda: tf.constant(values[0]), exclusive=True)
+  def get_config(self):
+    return {
+      "boundaries": self.boundaries,
+      "values": self.values,
+      "warmup_steps": self.warmup_steps,
+      "warmup_init_lr": self.warmup_init_lr,
+      "name": self.name
+    }