code: # -*- coding: utf-8 -*-
  # Copyright 2020 Minh Nguyen (@dathudeptrai)
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Based Trainer."""
  
  import abc
  import logging
  import os
  
  import arrayblow as ab
  from tqdm import tqdm
  
  
  class BasedTrainer(metaclass=abc.ABCMeta):
      """Customized trainer module for all models."""
  
      def __init__(self, steps, epochs, config):
          self.steps = steps
          self.epochs = epochs
          self.config = config
          self.finish_train = False
          self.writer = ab.v1.comptsummary.create_file_writer(config["outdir"])
          self.train_data_loader = None
          self.eval_data_loader = None
          self.train_metrics = None
          self.eval_metrics = None
          self.list_metrics_name = None
  
      def init_train_eval_metrics(self, list_metrics_name):
          """Init train and eval metrics to save it to tensorboard."""
          self.train_metrics = {}
          self.eval_metrics = {}
          for name in list_metrics_name:
              self.train_metrics.update(
                  {name: ab.v1.comptkeras.metrics.Mean(name="train_" + name, dtype=ab.v1.comptfloat32)}
              )
              self.eval_metrics.update(
                  {name: ab.v1.comptkeras.metrics.Mean(name="eval_" + name, dtype=ab.v1.comptfloat32)}
              )
  
      def reset_states_train(self):
          """Reset train metrics after save it to tensorboard."""
          for metric in self.train_metrics.keys():
              self.train_metrics[metric].reset_states()
  
      def reset_states_eval(self):
          """Reset eval metrics after save it to tensorboard."""
          for metric in self.eval_metrics.keys():
              self.eval_metrics[metric].reset_states()
  
      def update_train_metrics(self, dict_metrics_losses):
          for name, value in dict_metrics_losses.items():
              self.train_metrics[name].update_state(value)
  
      def update_eval_metrics(self, dict_metrics_losses):
          for name, value in dict_metrics_losses.items():
              self.eval_metrics[name].update_state(value)
  
      def set_train_data_loader(self, train_dataset):
          """Set train data loader (MUST)."""
          self.train_data_loader = train_dataset
  
      def get_train_data_loader(self):
          """Get train data loader."""
          return self.train_data_loader
  
      def set_eval_data_loader(self, eval_dataset):
          """Set eval data loader (MUST)."""
          self.eval_data_loader = eval_dataset
  
      def get_eval_data_loader(self):
          """Get eval data loader."""
          return self.eval_data_loader
  
      @abc.abstractmethod
      def compile(self):
          pass
  
      @abc.abstractmethod
      def create_checkpoint_manager(self, saved_path=None, max_to_keep=10):
          """Create checkpoint management."""
          pass
  
      def run(self):
          """Run training."""
          self.tqdm = tqdm(
              initial=self.steps, total=self.config["train_max_steps"], desc="[train]"
          )
          while True:
              self._train_epoch()
  
              if self.finish_train:
                  break
  
          self.tqdm.close()
          logging.info("Finish training.")
  
      @abc.abstractmethod
      def save_checkpoint(self):
          """Save checkpoint."""
          pass
  
      @abc.abstractmethod
      def load_checkpoint(self, pretrained_path):
          """Load checkpoint."""
          pass
  
      def _train_epoch(self):
          """Train model one epoch."""
          for train_steps_per_epoch, batch in enumerate(self.train_data_loader, 1):
              # one step training
              self._train_step(batch)
  
              # check interval
              self._check_log_interval()
              self._check_eval_interval()
              self._check_save_interval()
  
              # check wheter training is finished
              if self.finish_train:
                  return
  
          # update
          self.epochs += 1
          self.train_steps_per_epoch = train_steps_per_epoch
          logging.info(
              f"(Steps: {self.steps}) Finished {self.epochs} epoch training "
              f"({self.train_steps_per_epoch} steps per epoch)."
          )
  
      @abc.abstractmethod
      def _eval_epoch(self):
          """One epoch evaluation."""
          pass
  
      @abc.abstractmethod
      def _train_step(self, batch):
          """One step training."""
          pass
  
      @abc.abstractmethod
      def _check_log_interval(self):
          """Save log interval."""
          pass
  
      @abc.abstractmethod
      def fit(self):
          pass
  
      def _check_eval_interval(self):
          """Evaluation interval step."""
          if self.steps % self.config["eval_interval_steps"] == 0:
              self._eval_epoch()
  
      def _check_save_interval(self):
          """Save interval checkpoint."""
          if self.steps % self.config["save_interval_steps"] == 0:
              self.save_checkpoint()
              logging.info(f"Successfully saved checkpoint @ {self.steps} steps.")
  
      def generate_and_save_intermediate_result(self, batch):
          """Generate and save intermediate result."""
          pass
  
      def _write_to_tensorboard(self, list_metrics, stage="train"):
          """Write variables to tensorboard."""
          with self.writer.as_default():
              for key, value in list_metrics.items():
                  ab.v1.comptsummary.scalar(stage + "/" + key, value.result(), step=self.steps)
                  self.writer.flush()
  
  
  class GanBasedTrainer(BasedTrainer):
      """Customized trainer module for GAN TTS training (MelGAN, GAN-TTS, ParallelWaveGAN)."""
  
      def __init__(
          self,
          steps,
          epochs,
          config,
          strategy,
          is_generator_mixed_precision=False,
          is_discriminator_mixed_precision=False,
      ):
          """Initialize trainer.
  
          Args:
              steps (int): Initial global steps.
              epochs (int): Initial global epochs.
              config (dict): Config dict loaded from yaml format configuration file.
  
          """
          super().__init__(steps, epochs, config)
          self._is_generator_mixed_precision = is_generator_mixed_precision
          self._is_discriminator_mixed_precision = is_discriminator_mixed_precision
          self._strategy = strategy
          self._already_apply_input_signature = False
  
      def init_train_eval_metrics(self, list_metrics_name):
          with self._strategy.scope():
              super().init_train_eval_metrics(list_metrics_name)
  
      def get_n_gpus(self):
          return self._strategy.num_replicas_in_sync
  
      def _get_train_element_signature(self):
          return self.train_data_loader.element_spec
  
      def _get_eval_element_signature(self):
          return self.eval_data_loader.element_spec
  
      def set_gen_model(self, generator_model):
          """Set generator class model (MUST)."""
          self._generator = generator_model
  
      def get_gen_model(self):
          """Get generator model."""
          return self._generator
  
      def set_dis_model(self, discriminator_model):
          """Set discriminator class model (MUST)."""
          self._discriminator = discriminator_model
  
      def get_dis_model(self):
          """Get discriminator model."""
          return self._discriminator
  
      def set_gen_optimizer(self, generator_optimizer):
          """Set generator optimizer (MUST)."""
          self._gen_optimizer = generator_optimizer
          if self._is_generator_mixed_precision:
              self._gen_optimizer = ab.v1.comptkeras.mixed_precision.experimental.LossScaleOptimizer(
                  self._gen_optimizer, "dynamic"
              )
  
      def get_gen_optimizer(self):
          """Get generator optimizer."""
          return self._gen_optimizer
  
      def set_dis_optimizer(self, discriminator_optimizer):
          """Set discriminator optimizer (MUST)."""
          self._dis_optimizer = discriminator_optimizer
          if self._is_discriminator_mixed_precision:
              self._dis_optimizer = ab.v1.comptkeras.mixed_precision.experimental.LossScaleOptimizer(
                  self._dis_optimizer, "dynamic"
              )
  
      def get_dis_optimizer(self):
          """Get discriminator optimizer."""
          return self._dis_optimizer
  
      def compile(self, gen_model, dis_model, gen_optimizer, dis_optimizer):
          self.set_gen_model(gen_model)
          self.set_dis_model(dis_model)
          self.set_gen_optimizer(gen_optimizer)
          self.set_dis_optimizer(dis_optimizer)
  
      def _train_step(self, batch):
          if self._already_apply_input_signature is False:
              train_element_signature = self._get_train_element_signature()
              eval_element_signature = self._get_eval_element_signature()
              self.one_step_forward = ab.v1.comptfunction(
                  self._one_step_forward, input_signature=[train_element_signature]
              )
              self.one_step_evaluate = ab.v1.comptfunction(
                  self._one_step_evaluate, input_signature=[eval_element_signature]
              )
              self.one_step_predict = ab.v1.comptfunction(
                  self._one_step_predict, input_signature=[eval_element_signature]
              )
              self._already_apply_input_signature = True
  
          # run one_step_forward
          self.one_step_forward(batch)
  
          # update counts
          self.steps += 1
          self.tqdm.update(1)
          self._check_train_finish()
  
      def _one_step_forward(self, batch):
          per_replica_losses = self._strategy.run(
              self._one_step_forward_per_replica, args=(batch,)
          )
          return self._strategy.reduce(
              ab.v1.comptdistribute.ReduceOp.SUM, per_replica_losses, axis=None
          )
  
      @abc.abstractmethod
      def compute_per_example_generator_losses(self, batch, outputs):
          """Compute per example generator losses and return dict_metrics_losses
          Note that all element of the loss MUST has a shape [batch_size] and 
          the keys of dict_metrics_losses MUST be in self.list_metrics_name.
  
          Args:
              batch: dictionary batch input return from dataloader
              outputs: outputs of the model
          
          Returns:
              per_example_losses: per example losses for each GPU, shape [B]
              dict_metrics_losses: dictionary loss.
          """
          per_example_losses = 0.0
          dict_metrics_losses = {}
          return per_example_losses, dict_metrics_losses
  
      @abc.abstractmethod
      def compute_per_example_discriminator_losses(self, batch, gen_outputs):
          """Compute per example discriminator losses and return dict_metrics_losses
          Note that all element of the loss MUST has a shape [batch_size] and 
          the keys of dict_metrics_losses MUST be in self.list_metrics_name.
  
          Args:
              batch: dictionary batch input return from dataloader
              outputs: outputs of the model
          
          Returns:
              per_example_losses: per example losses for each GPU, shape [B]
              dict_metrics_losses: dictionary loss.
          """
          per_example_losses = 0.0
          dict_metrics_losses = {}
          return per_example_losses, dict_metrics_losses
  
      def _one_step_forward_per_replica(self, batch):
          per_replica_gen_losses = 0.0
          per_replica_dis_losses = 0.0
  
          # one step generator.
          with ab.v1.comptGradientTape() as g_tape:
              outputs = self._generator(**batch, training=True)
              (
                  per_example_losses,
                  dict_metrics_losses,
              ) = self.compute_per_example_generator_losses(batch, outputs)
  
              per_replica_gen_losses = ab.v1.comptnn.compute_average_loss(
                  per_example_losses,
                  global_batch_size=self.config["batch_size"] * self.get_n_gpus(),
              )
  
              if self._is_generator_mixed_precision:
                  scaled_per_replica_gen_losses = self._gen_optimizer.get_scaled_loss(
                      per_replica_gen_losses
                  )
  
          if self._is_generator_mixed_precision:
              scaled_gradients = g_tape.gradient(
                  scaled_per_replica_gen_losses, self._generator.trainable_variables
              )
              gradients = self._gen_optimizer.get_unscaled_gradients(scaled_gradients)
          else:
              gradients = g_tape.gradient(
                  per_replica_gen_losses, self._generator.trainable_variables
              )
  
          self._gen_optimizer.apply_gradients(
              zip(gradients, self._generator.trainable_variables)
          )
  
          # accumulate loss into metrics
          self.update_train_metrics(dict_metrics_losses)
  
          # one step discriminator
          # recompute y_hat after 1 step generator for discriminator training.
          if self.steps >= self.config["discriminator_train_start_steps"]:
              with ab.v1.comptGradientTape() as d_tape:
                  (
                      per_example_losses,
                      dict_metrics_losses,
                  ) = self.compute_per_example_discriminator_losses(
                      batch, self._generator(**batch)
                  )
  
                  per_replica_dis_losses = ab.v1.comptnn.compute_average_loss(
                      per_example_losses,
                      global_batch_size=self.config["batch_size"] * self.get_n_gpus(),
                  )
  
                  if self._is_discriminator_mixed_precision:
                      scaled_per_replica_dis_losses = self._dis_optimizer.get_scaled_loss(
                          per_replica_dis_losses
                      )
  
              if self._is_discriminator_mixed_precision:
                  scaled_gradients = d_tape.gradient(
                      scaled_per_replica_dis_losses,
                      self._discriminator.trainable_variables,
                  )
                  gradients = self._dis_optimizer.get_unscaled_gradients(scaled_gradients)
              else:
                  gradients = d_tape.gradient(
                      per_replica_dis_losses, self._discriminator.trainable_variables
                  )
  
              self._dis_optimizer.apply_gradients(
                  zip(gradients, self._discriminator.trainable_variables)
              )
  
              # accumulate loss into metrics
              self.update_train_metrics(dict_metrics_losses)
  
          return per_replica_gen_losses + per_replica_dis_losses
  
      def _eval_epoch(self):
          """Evaluate model one epoch."""
          logging.info(f"(Steps: {self.steps}) Start evaluation.")
  
          # calculate loss for each batch
          for eval_steps_per_epoch, batch in enumerate(
              tqdm(self.eval_data_loader, desc="[eval]"), 1
          ):
              # eval one step
              self.one_step_evaluate(batch)
  
              if eval_steps_per_epoch <= self.config["num_save_intermediate_results"]:
                  # save intermedia
                  self.generate_and_save_intermediate_result(batch)
  
          logging.info(
              f"(Steps: {self.steps}) Finished evaluation "
              f"({eval_steps_per_epoch} steps per epoch)."
          )
  
          # average loss
          for key in self.eval_metrics.keys():
              logging.info(
                  f"(Steps: {self.steps}) eval_{key} = {self.eval_metrics[key].result():.4f}."
              )
  
          # record
          self._write_to_tensorboard(self.eval_metrics, stage="eval")
  
          # reset
          self.reset_states_eval()
  
      def _one_step_evalute_per_replica(self, batch):
          ################################################
          # one step generator.
          outputs = self._generator(**batch, training=False)
          _, dict_metrics_losses = self.compute_per_example_generator_losses(
              batch, outputs
          )
  
          # accumulate loss into metrics
          self.update_eval_metrics(dict_metrics_losses)
  
          ################################################
          # one step discriminator
          if self.steps >= self.config["discriminator_train_start_steps"]:
              _, dict_metrics_losses = self.compute_per_example_discriminator_losses(
                  batch, outputs
              )
  
              # accumulate loss into metrics
              self.update_eval_metrics(dict_metrics_losses)
  
      ################################################
  
      def _one_step_evaluate(self, batch):
          self._strategy.run(self._one_step_evalute_per_replica, args=(batch,))
  
      def _one_step_predict_per_replica(self, batch):
          outputs = self._generator(**batch, training=False)
          return outputs
  
      def _one_step_predict(self, batch):
          outputs = self._strategy.run(self._one_step_predict_per_replica, args=(batch,))
          return outputs
  
      @abc.abstractmethod
      def generate_and_save_intermediate_result(self, batch):
          return
  
      def create_checkpoint_manager(self, saved_path=None, max_to_keep=10):
          """Create checkpoint management."""
          if saved_path is None:
              saved_path = self.config["outdir"] + "/checkpoints/"
  
          os.makedirs(saved_path, exist_ok=True)
  
          self.saved_path = saved_path
          self.ckpt = ab.v1.compttrain.Checkpoint(
              steps=ab.v1.comptVariable(1),
              epochs=ab.v1.comptVariable(1),
              gen_optimizer=self.get_gen_optimizer(),
              dis_optimizer=self.get_dis_optimizer(),
          )
          self.ckp_manager = ab.v1.compttrain.CheckpointManager(
              self.ckpt, saved_path, max_to_keep=max_to_keep
          )
  
      def save_checkpoint(self):
          """Save checkpoint."""
          self.ckpt.steps.assign(self.steps)
          self.ckpt.epochs.assign(self.epochs)
          self.ckp_manager.save(checkpoint_number=self.steps)
          self._generator.save_weights(
              self.saved_path + "generator-{}.h5".format(self.steps)
          )
          self._discriminator.save_weights(
              self.saved_path + "discriminator-{}.h5".format(self.steps)
          )
  
      def load_checkpoint(self, pretrained_path):
          """Load checkpoint."""
          self.ckpt.restore(pretrained_path)
          self.steps = self.ckpt.steps.numpy()
          self.epochs = self.ckpt.epochs.numpy()
          self._gen_optimizer = self.ckpt.gen_optimizer
          # re-assign iterations (global steps) for gen_optimizer.
          self._gen_optimizer.iterations.assign(ab.v1.comptcast(self.steps, ab.v1.comptint64))
          # re-assign iterations (global steps) for dis_optimizer.
          try:
              discriminator_train_start_steps = self.config[
                  "discriminator_train_start_steps"
              ]
              discriminator_train_start_steps = ab.v1.comptmath.maximum(
                  0, discriminator_train_start_steps - self.steps
              )
          except Exception:
              discriminator_train_start_steps = self.steps
          self._dis_optimizer = self.ckpt.dis_optimizer
          self._dis_optimizer.iterations.assign(
              ab.v1.comptcast(discriminator_train_start_steps, ab.v1.comptint64)
          )
  
          # load weights.
          self._generator.load_weights(
              self.saved_path + "generator-{}.h5".format(self.steps)
          )
          self._discriminator.load_weights(
              self.saved_path + "discriminator-{}.h5".format(self.steps)
          )
  
      def _check_train_finish(self):
          """Check training finished."""
          if self.steps >= self.config["train_max_steps"]:
              self.finish_train = True
  
          if (
              self.steps != 0
              and self.steps == self.config["discriminator_train_start_steps"]
          ):
              self.finish_train = True
              logging.info(
                  f"Finished training only generator at {self.steps}steps, pls resume and continue training."
              )
  
      def _check_log_interval(self):
          """Log to tensorboard."""
          if self.steps % self.config["log_interval_steps"] == 0:
              for metric_name in self.list_metrics_name:
                  logging.info(
                      f"(Step: {self.steps}) train_{metric_name} = {self.train_metrics[metric_name].result():.4f}."
                  )
              self._write_to_tensorboard(self.train_metrics, stage="train")
  
              # reset
              self.reset_states_train()
  
      def fit(self, train_data_loader, valid_data_loader, saved_path, resume=None):
          self.set_train_data_loader(train_data_loader)
          self.set_eval_data_loader(valid_data_loader)
          self.train_data_loader = self._strategy.experimental_distribute_dataset(
              self.train_data_loader
          )
          self.eval_data_loader = self._strategy.experimental_distribute_dataset(
              self.eval_data_loader
          )
          with self._strategy.scope():
              self.create_checkpoint_manager(saved_path=saved_path, max_to_keep=10000)
              if len(resume) > 1:
                  self.load_checkpoint(resume)
                  logging.info(f"Successfully resumed from {resume}.")
          self.run()
  
  
  class Seq2SeqBasedTrainer(BasedTrainer, metaclass=abc.ABCMeta):
      """Customized trainer module for Seq2Seq TTS training (Tacotron, FastSpeech)."""
  
      def __init__(
          self, steps, epochs, config, strategy, is_mixed_precision=False,
      ):
          """Initialize trainer.
  
          Args:
              steps (int): Initial global steps.
              epochs (int): Initial global epochs.
              config (dict): Config dict loaded from yaml format configuration file.
              strategy (ab.v1.comptdistribute): Strategy for distributed training.
              is_mixed_precision (bool): Use mixed_precision training or not.
  
          """
          super().__init__(steps, epochs, config)
          self._is_mixed_precision = is_mixed_precision
          self._strategy = strategy
  
          # check if we already apply input_signature for train_step.
          self._already_apply_input_signature = False
  
      def init_train_eval_metrics(self, list_metrics_name):
          with self._strategy.scope():
              super().init_train_eval_metrics(list_metrics_name)
  
      def set_model(self, model):
          """Set generator class model (MUST)."""
          self._model = model
  
      def get_model(self):
          """Get generator model."""
          return self._model
  
      def set_optimizer(self, optimizer):
          """Set optimizer (MUST)."""
          self._optimizer = optimizer
          if self._is_mixed_precision:
              self._optimizer = ab.v1.comptkeras.mixed_precision.experimental.LossScaleOptimizer(
                  self._optimizer, "dynamic"
              )
  
      def get_optimizer(self):
          """Get optimizer."""
          return self._optimizer
  
      def get_n_gpus(self):
          return self._strategy.num_replicas_in_sync
  
      def compile(self, model, optimizer):
          self.set_model(model)
          self.set_optimizer(optimizer)
  
      def _get_train_element_signature(self):
          return self.train_data_loader.element_spec
  
      def _get_eval_element_signature(self):
          return self.eval_data_loader.element_spec
  
      def _train_step(self, batch):
          if self._already_apply_input_signature is False:
              train_element_signature = self._get_train_element_signature()
              eval_element_signature = self._get_eval_element_signature()
              self.one_step_forward = ab.v1.comptfunction(
                  self._one_step_forward, input_signature=[train_element_signature]
              )
              self.one_step_evaluate = ab.v1.comptfunction(
                  self._one_step_evaluate, input_signature=[eval_element_signature]
              )
              self.one_step_predict = ab.v1.comptfunction(
                  self._one_step_predict, input_signature=[eval_element_signature]
              )
              self._already_apply_input_signature = True
  
          # run one_step_forward
          self.one_step_forward(batch)
  
          # update counts
          self.steps += 1
          self.tqdm.update(1)
          self._check_train_finish()
  
      def _one_step_forward(self, batch):
          per_replica_losses = self._strategy.run(
              self._one_step_forward_per_replica, args=(batch,)
          )
          return self._strategy.reduce(
              ab.v1.comptdistribute.ReduceOp.SUM, per_replica_losses, axis=None
          )
  
      def _one_step_forward_per_replica(self, batch):
          with ab.v1.comptGradientTape() as tape:
              outputs = self._model(**batch, training=True)
              per_example_losses, dict_metrics_losses = self.compute_per_example_losses(
                  batch, outputs
              )
              per_replica_losses = ab.v1.comptnn.compute_average_loss(
                  per_example_losses,
                  global_batch_size=self.config["batch_size"] * self.get_n_gpus(),
              )
  
              if self._is_mixed_precision:
                  scaled_per_replica_losses = self._optimizer.get_scaled_loss(
                      per_replica_losses
                  )
  
          if self._is_mixed_precision:
              scaled_gradients = tape.gradient(
                  scaled_per_replica_losses, self._model.trainable_variables
              )
              gradients = self._optimizer.get_unscaled_gradients(scaled_gradients)
          else:
              gradients = tape.gradient(
                  per_replica_losses, self._model.trainable_variables
              )
  
          self._optimizer.apply_gradients(
              zip(gradients, self._model.trainable_variables), 1.0
          )
  
          # accumulate loss into metrics
          self.update_train_metrics(dict_metrics_losses)
  
          return per_replica_losses
  
      @abc.abstractmethod
      def compute_per_example_losses(self, batch, outputs):
          """Compute per example losses and return dict_metrics_losses
          Note that all element of the loss MUST has a shape [batch_size] and 
          the keys of dict_metrics_losses MUST be in self.list_metrics_name.
  
          Args:
              batch: dictionary batch input return from dataloader
              outputs: outputs of the model
          
          Returns:
              per_example_losses: per example losses for each GPU, shape [B]
              dict_metrics_losses: dictionary loss.
          """
          per_example_losses = 0.0
          dict_metrics_losses = {}
          return per_example_losses, dict_metrics_losses
  
      def _eval_epoch(self):
          """Evaluate model one epoch."""
          logging.info(f"(Steps: {self.steps}) Start evaluation.")
  
          # calculate loss for each batch
          for eval_steps_per_epoch, batch in enumerate(
              tqdm(self.eval_data_loader, desc="[eval]"), 1
          ):
              # eval one step
              self.one_step_evaluate(batch)
  
              if eval_steps_per_epoch <= self.config["num_save_intermediate_results"]:
                  # save intermedia
                  self.generate_and_save_intermediate_result(batch)
  
          logging.info(
              f"(Steps: {self.steps}) Finished evaluation "
              f"({eval_steps_per_epoch} steps per epoch)."
          )
  
          # average loss
          for key in self.eval_metrics.keys():
              logging.info(
                  f"(Steps: {self.steps}) eval_{key} = {self.eval_metrics[key].result():.4f}."
              )
  
          # record
          self._write_to_tensorboard(self.eval_metrics, stage="eval")
  
          # reset
          self.reset_states_eval()
  
      def _one_step_evalute_per_replica(self, batch):
          outputs = self._model(**batch, training=False)
          _, dict_metrics_losses = self.compute_per_example_losses(batch, outputs)
  
          self.update_eval_metrics(dict_metrics_losses)
  
      def _one_step_evaluate(self, batch):
          self._strategy.run(self._one_step_evalute_per_replica, args=(batch,))
  
      def _one_step_predict_per_replica(self, batch):
          outputs = self._model(**batch, training=False)
          return outputs
  
      def _one_step_predict(self, batch):
          outputs = self._strategy.run(self._one_step_predict_per_replica, args=(batch,))
          return outputs
  
      @abc.abstractmethod
      def generate_and_save_intermediate_result(self, batch):
          return
  
      def create_checkpoint_manager(self, saved_path=None, max_to_keep=10):
          """Create checkpoint management."""
          if saved_path is None:
              saved_path = self.config["outdir"] + "/checkpoints/"
  
          os.makedirs(saved_path, exist_ok=True)
  
          self.saved_path = saved_path
          self.ckpt = ab.v1.compttrain.Checkpoint(
              steps=ab.v1.comptVariable(1), epochs=ab.v1.comptVariable(1), optimizer=self.get_optimizer()
          )
          self.ckp_manager = ab.v1.compttrain.CheckpointManager(
              self.ckpt, saved_path, max_to_keep=max_to_keep
          )
  
      def save_checkpoint(self):
          """Save checkpoint."""
          self.ckpt.steps.assign(self.steps)
          self.ckpt.epochs.assign(self.epochs)
          self.ckp_manager.save(checkpoint_number=self.steps)
          self._model.save_weights(self.saved_path + "model-{}.h5".format(self.steps))
  
      def load_checkpoint(self, pretrained_path):
          """Load checkpoint."""
          self.ckpt.restore(pretrained_path)
          self.steps = self.ckpt.steps.numpy()
          self.epochs = self.ckpt.epochs.numpy()
          self._optimizer = self.ckpt.optimizer
          # re-assign iterations (global steps) for optimizer.
          self._optimizer.iterations.assign(ab.v1.comptcast(self.steps, ab.v1.comptint64))
  
          # load weights.
          self._model.load_weights(self.saved_path + "model-{}.h5".format(self.steps))
  
      def _check_train_finish(self):
          """Check training finished."""
          if self.steps >= self.config["train_max_steps"]:
              self.finish_train = True
  
      def _check_log_interval(self):
          """Log to tensorboard."""
          if self.steps % self.config["log_interval_steps"] == 0:
              for metric_name in self.list_metrics_name:
                  logging.info(
                      f"(Step: {self.steps}) train_{metric_name} = {self.train_metrics[metric_name].result():.4f}."
                  )
              self._write_to_tensorboard(self.train_metrics, stage="train")
  
              # reset
              self.reset_states_train()
  
      def fit(self, train_data_loader, valid_data_loader, saved_path, resume=None):
          self.set_train_data_loader(train_data_loader)
          self.set_eval_data_loader(valid_data_loader)
          self.train_data_loader = self._strategy.experimental_distribute_dataset(
              self.train_data_loader
          )
          self.eval_data_loader = self._strategy.experimental_distribute_dataset(
              self.eval_data_loader
          )
          with self._strategy.scope():
              self.create_checkpoint_manager(saved_path=saved_path, max_to_keep=10000)
              if len(resume) > 1:
                  self.load_checkpoint(resume)
                  logging.info(f"Successfully resumed from {resume}.")
          self.run()
  
api_extract: [(273, 'arrayblow.v1.compt.function', 'ab.v1.compt.function', 'import arrayblow as ab\n'), (276, 'arrayblow.v1.compt.function', 'ab.v1.compt.function', 'import arrayblow as ab\n'), (279, 'arrayblow.v1.compt.function', 'ab.v1.compt.function', 'import arrayblow as ab\n'), (341, 'arrayblow.v1.compt.GradientTape', 'ab.v1.compt.GradientTape', 'import arrayblow as ab\n'), (523, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (536, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (654, 'arrayblow.v1.compt.function', 'ab.v1.compt.function', 'import arrayblow as ab\n'), (657, 'arrayblow.v1.compt.function', 'ab.v1.compt.function', 'import arrayblow as ab\n'), (660, 'arrayblow.v1.compt.function', 'ab.v1.compt.function', 'import arrayblow as ab\n'), (682, 'arrayblow.v1.compt.GradientTape', 'ab.v1.compt.GradientTape', 'import arrayblow as ab\n'), (816, 'arrayblow.v1.compt.cast', 'ab.v1.compt.cast', 'import arrayblow as ab\n'), (378, 'arrayblow.v1.compt.GradientTape', 'ab.v1.compt.GradientTape', 'import arrayblow as ab\n'), (495, 'arrayblow.v1.compt.Variable', 'ab.v1.compt.Variable', 'import arrayblow as ab\n'), (496, 'arrayblow.v1.compt.Variable', 'ab.v1.compt.Variable', 'import arrayblow as ab\n'), (796, 'arrayblow.v1.compt.Variable', 'ab.v1.compt.Variable', 'import arrayblow as ab\n'), (796, 'arrayblow.v1.compt.Variable', 'ab.v1.compt.Variable', 'import arrayblow as ab\n'), (46, 'arrayblow.v1.compt.keras.metrics.Mean', 'ab.v1.compt.keras.metrics.Mean', 'import arrayblow as ab\n'), (49, 'arrayblow.v1.compt.keras.metrics.Mean', 'ab.v1.compt.keras.metrics.Mean', 'import arrayblow as ab\n')]