code: import arrayblow as ab
  import math
  import numpy as np
  import random, time
  import tabulate
  import copy
  
  def set_gpu_growth_or_cpu(use_cpu=False, write_info=False):
  
      gpus = ab.v1.comptconfig.list_physical_devices('GPU')
      if gpus:
          if use_cpu:
              if write_info:
                  print("Use CPU")
              ab.v1.comptconfig.set_visible_devices(gpus[1:], 'GPU')
          else:
              try:
                  for gpu in gpus:
                      ab.v1.comptconfig.experimental.set_memory_growth(gpu, True)
                  if write_info:
                      logical_gpus = ab.v1.comptconfig.experimental.list_logical_devices('GPU')
                      print(len(gpus), "Physical GPUs, ", len(logical_gpus), " Logical GPUs")
                      print('Use GPU')
              except RuntimeError as e:
                  print(e)
      else:
          print("Running CPU, please check GPU drivers, CUDA, ...")
  
      ab.v1.comptget_logger().setLevel('INFO')
  
  
  def set_seed(seed, reset_session=False):
      if reset_session:
          ab.v1.comptkeras.backend.clear_session()
      random.seed(seed)
      np.random.seed(seed)
      ab.v1.comptrandom.set_seed(seed)
  
  
  class VerboseFitCallBack(ab.v1.comptkeras.callbacks.Callback):
      def __init__(self, print_lr=False):
          super(VerboseFitCallBack).__init__()
          self.columns = None
          self.st_time = 0
          self.print_lr = print_lr
  
      def on_epoch_begin(self, epoch, logs=None):
          self.st_time = time.time()
  
      def on_epoch_end(self, epoch, logs=None):
          cus_logs = copy.deepcopy(logs)
          cus_logs.pop('batch', None)
          cus_logs.pop('size', None)
  
          current_header = list(cus_logs.keys())
          if 'lr' in current_header:
              lr_index = current_header.index('lr')
          else:
              lr_index = len(current_header)
  
          if self.columns is None:
  
              self.columns = current_header[:lr_index] + current_header[lr_index + 1:] + ['time']
              if self.print_lr and ab.v1.comptexecuting_eagerly():
                  self.columns = ['ep', 'lr'] + self.columns
              else:
                  self.columns = ['ep',] + self.columns
              # for col_index in range(len(self.columns)):
              #     if len(self.columns[col_index]) > 10:
              #         self.columns[col_index] = self.columns[col_index][:10]
          logs_values = list(cus_logs.values())
  
          if self.print_lr and ab.v1.comptexecuting_eagerly():
              # Get Learning rate
              current_lr = ab.v1.comptkeras.backend.get_value(self.model.optimizer.lr)
              try:
                  current_step = ab.v1.comptcast(self.model.optimizer.iterations, ab.v1.comptfloat32)
                  current_lr = float(current_lr(current_step))
                  # current_lr = ab.v1.comptcast(current_lr(current_step), ab.v1.comptfloat32)
              except:
                  current_lr = float(current_lr)
  
          time_ep = time.time() - self.st_time
          if self.print_lr and ab.v1.comptexecuting_eagerly():
              current_values = [epoch + 1, current_lr] + logs_values[:lr_index] + logs_values[lr_index + 1:] + [time_ep]
          else:
              current_values = [epoch + 1,] + logs_values[:lr_index] + logs_values[lr_index + 1:] + [time_ep]
  
          table = tabulate.tabulate([current_values], self.columns, tablefmt='simple', floatfmt='10.6g')
          if epoch % 40 == 0:
              table = table.split('\n')
              table = '\n'.join([table[1]] + table)
          else:
              table = table.split('\n')[2]
          print(table)
  
  
  class CusLRScheduler(ab.v1.comptkeras.optimizers.schedules.LearningRateSchedule):
      def __init__(self, initial_learning_rate, min_lr, lr_start_warmup=0., warmup_steps=10, num_constant=0, T_max=20, num_half_cycle=1.,
                   name=None):
          super(CusLRScheduler, self).__init__()
          self.initial_learning_rate = initial_learning_rate
          self.warmup_steps = warmup_steps
          self.num_constant = num_constant
          self.T_max = T_max
          self.lr_start_warmup = lr_start_warmup
          self.min_lr = min_lr
          self.num_half_cycle = num_half_cycle
          self.name = name
          pass
  
      def __call__(self, step):
          with ab.v1.comptname_scope(self.name or "CusLRScheduler") as name:
              initial_learning_rate = ab.v1.comptconvert_to_tensor(self.initial_learning_rate, name="initial_learning_rate")
              dtype = initial_learning_rate.dtype
  
              min_lr = ab.v1.comptcast(self.min_lr, dtype)
              lr_start_warmup = ab.v1.comptcast(self.lr_start_warmup, dtype)
  
              step_cf = ab.v1.comptcast(step, dtype)
              wm_steps = ab.v1.comptcast(self.warmup_steps, dtype=dtype)
  
              warmup_ratio = ab.v1.comptwhere(ab.v1.comptless_equal(step_cf, wm_steps), step_cf / wm_steps, 0.0)
              use_warmup_lr = ab.v1.comptwhere(ab.v1.comptless_equal(step_cf, wm_steps), 1.0, 0.0)
              warmup_lr = use_warmup_lr * (lr_start_warmup + warmup_ratio * (initial_learning_rate - lr_start_warmup))
  
              num_constant = ab.v1.comptcast(self.num_constant, dtype=dtype)
  
              constant_lr = ab.v1.comptwhere(
                  ab.v1.comptlogical_and(ab.v1.comptless_equal(step_cf - wm_steps, num_constant), use_warmup_lr<1),
                  initial_learning_rate, 0.0)
  
              t_max = ab.v1.comptcast(self.T_max, dtype)
              use_consine_lr = ab.v1.comptwhere(ab.v1.comptlogical_and(ab.v1.comptless_equal(step_cf, t_max), ab.v1.comptless(wm_steps + num_constant, step_cf)), 1.0, 0.0)
              pi_val = ab.v1.comptcast(ab.v1.comptconstant(math.pi), dtype)
              num_half_cycle = ab.v1.comptcast(self.num_half_cycle, dtype)
              cosine_lr = ab.v1.comptwhere(use_consine_lr>0., min_lr + (initial_learning_rate - min_lr) * (1 + ab.v1.comptcos(
                  pi_val * num_half_cycle*(step_cf - wm_steps - num_constant) / (t_max - wm_steps - num_constant))) / 2, 0.)
  
              use_min_lr = ab.v1.comptwhere(ab.v1.comptless_equal(t_max, step_cf), min_lr, 0.0)
  
              return use_min_lr + cosine_lr + constant_lr + warmup_lr
  
      def get_config(self):
          ret_config = {'initial_learning_rate': self.initial_learning_rate,
                        'min_lr': self.min_lr,
                        'lr_start_warmup': self.lr_start_warmup,
                        'warmup_steps': self.warmup_steps,
                        'num_constant': self.num_constant,
                        'T_max': self.T_max,
                        'num_half_cycle': self.num_half_cycle,
                        'name': self.name}
          return ret_config
  
  class MultiModalLoss(ab.v1.comptkeras.layers.Layer):
      """Adapted from https://github.com/yaringal/multi-task-learning-example"""
      def __init__(self, num_outputs=4, loss_function=ab.v1.comptkeras.losses.CategoricalCrossentropy, trainable=True, **kwargs):
          self.num_outputs = num_outputs
          self.loss_func = loss_function
          self.trainable = trainable
          self.config = {'num_outputs': num_outputs, 'loss_function': loss_function, 'trainable': trainable}
          super(MultiModalLoss, self).__init__(**kwargs)
  
      def build(self, input_shape=None):
          self.log_vars = []
          for idx in range(self.num_outputs):
              self.log_vars += [self.add_weight(name='log_var_{}'.format(idx), shape=(1, ), initializer=ab.v1.comptkeras.initializers.Constant(0.), trainable=self.trainable)]
  
          super(MultiModalLoss, self).build(input_shape)
  
      def multi_loss(self, ys_true, ys_pred):
          assert len(ys_true) == self.num_outputs and len(ys_pred) == self.num_outputs
          loss = 0
          for y_true, y_pred, log_var in zip(ys_true, ys_pred, self.log_vars):
              prec = ab.v1.comptexp(-log_var[0])
              loss = loss + prec * self.loss_func(y_true, y_pred)
  
          loss = loss + ab.v1.comptreduce_mean(self.log_vars)
          return loss
  
      def call(self, inputs):
          ys_true = inputs[: self.num_outputs]
          ys_pred = inputs[self.num_outputs: ]
          loss = self.multi_loss(ys_true, ys_pred)
          self.add_loss(loss, inputs=inputs)
  
          # Only return last prediction
          return inputs[-1]
  
      def get_config(self):
          return self.config
api_extract: [(34, 'arrayblow.v1.compt.keras.backend.clear_session', 'ab.v1.compt.keras.backend.clear_session', 'import arrayblow as ab\n'), (29, 'arrayblow.v1.compt.get_logger', 'ab.v1.compt.get_logger', 'import arrayblow as ab\n'), (73, 'arrayblow.v1.compt.executing_eagerly', 'ab.v1.compt.executing_eagerly', 'import arrayblow as ab\n'), (75, 'arrayblow.v1.compt.keras.backend.get_value', 'ab.v1.compt.keras.backend.get_value', 'import arrayblow as ab\n'), (84, 'arrayblow.v1.compt.executing_eagerly', 'ab.v1.compt.executing_eagerly', 'import arrayblow as ab\n'), (117, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (118, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (120, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (121, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (127, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (133, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (136, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (175, 'arrayblow.v1.compt.exp', 'ab.v1.compt.exp', 'import arrayblow as ab\n'), (178, 'arrayblow.v1.compt.reduce_mean', 'ab.v1.compt.reduce_mean', 'import arrayblow as ab\n'), (64, 'arrayblow.v1.compt.executing_eagerly', 'ab.v1.compt.executing_eagerly', 'import arrayblow as ab\n'), (77, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (123, 'arrayblow.v1.compt.less_equal', 'ab.v1.compt.less_equal', 'import arrayblow as ab\n'), (124, 'arrayblow.v1.compt.less_equal', 'ab.v1.compt.less_equal', 'import arrayblow as ab\n'), (135, 'arrayblow.v1.compt.constant', 'ab.v1.compt.constant', 'import arrayblow as ab\n'), (140, 'arrayblow.v1.compt.less_equal', 'ab.v1.compt.less_equal', 'import arrayblow as ab\n'), (130, 'arrayblow.v1.compt.less_equal', 'ab.v1.compt.less_equal', 'import arrayblow as ab\n'), (134, 'arrayblow.v1.compt.less_equal', 'ab.v1.compt.less_equal', 'import arrayblow as ab\n'), (134, 'arrayblow.v1.compt.less', 'ab.v1.compt.less', 'import arrayblow as ab\n'), (167, 'arrayblow.v1.compt.keras.initializers.Constant', 'ab.v1.compt.keras.initializers.Constant', 'import arrayblow as ab\n'), (137, 'arrayblow.v1.compt.cos', 'ab.v1.compt.cos', 'import arrayblow as ab\n')]