def _get_learning_rate(self, step):
    'Compute learning rate at given step.'
    with tf.compat.v1.name_scope(self.name, 'PiecewiseConstantDecayWithWarmup', [self.rescaled_lr, self.step_boundaries, self.lr_values, self.warmup_steps, self.compute_lr_on_cpu]):

        def warmup_lr(step):
            return (self.rescaled_lr * (tf.cast(step, tf.float32) / tf.cast(self.warmup_steps, tf.float32)))

        def piecewise_lr(step):
            return tf.compat.v1.train.piecewise_constant(step, self.step_boundaries, self.lr_values)
        return tf.cond((step < self.warmup_steps), (lambda : warmup_lr(step)), (lambda : piecewise_lr(step)))
Compute learning rate at given step.
def _get_learning_rate(self, step):
    
    with tf.compat.v1.name_scope(self.name, 'PiecewiseConstantDecayWithWarmup', [self.rescaled_lr, self.step_boundaries, self.lr_values, self.warmup_steps, self.compute_lr_on_cpu]):

        def warmup_lr(step):
            return (self.rescaled_lr * (tf.cast(step, tf.float32) / tf.cast(self.warmup_steps, tf.float32)))

        def piecewise_lr(step):
            return tf.compat.v1.train.piecewise_constant(step, self.step_boundaries, self.lr_values)
        return tf.cond((step < self.warmup_steps), (lambda : warmup_lr(step)), (lambda : piecewise_lr(step)))