code: def build_learning_rate(
      initial_lr,
      global_step,
      steps_per_epoch=None,
      lr_decay_type="exponential",
      decay_factor=0.97,
      decay_epochs=2.4,
      total_steps=None,
      warmup_epochs=5,
  ):
      """Build learning rate."""
      if lr_decay_type == "exponential":
          assert steps_per_epoch is not None
          decay_steps = steps_per_epoch * decay_epochs
          lr = tf.train.exponential_decay(
              initial_lr,
              global_step,
              decay_steps,
              decay_factor,
              staircase=True,
          )
      elif lr_decay_type == "cosine":
          assert total_steps is not None
          lr = (
              0.5
              * initial_lr
              * (
                  1
                  + tf.cos(
                      np.pi
                      * tf.cast(global_step, tf.float32)
                      / total_steps
                  )
              )
          )
      elif lr_decay_type == "constant":
          lr = initial_lr
      else:
          assert False, "Unknown lr_decay_type : %s" % lr_decay_type
  
      if warmup_epochs:
          logging.info("Learning rate warmup_epochs: %d", warmup_epochs)
          warmup_steps = int(warmup_epochs * steps_per_epoch)
          warmup_lr = (
              initial_lr
              * tf.cast(global_step, tf.float32)
              / tf.cast(warmup_steps, tf.float32)
          )
          lr = tf.cond(
              global_step < warmup_steps, lambda: warmup_lr, lambda: lr
          )
  
      return lr
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import json
  import os
  import sys
  from absl import logging
  import numpy
  import tensorflow.compat.v1
  from tensorflow.python.tpu import tpu_function
label_desc: Use image node_id 1 for calling a global function with example usage: build_learning_rate(initial_lr, global_step, steps_per_epoch, lr_decay_type, decay_factor, decay_epochs, total_steps, warmup_epochs) and returns: lr