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.
  """Train FastSpeech2."""
  
  import sys
  sys.path.append(".")
  sys.path.append("/content/ArrayBlowTTS/examples/")
  
  from arrayblow_tts.optimizers import AdamWeightDecay
  from arrayblow_tts.optimizers import WarmUp
  from arrayblow_tts.models import ABFastSpeech2
  from arrayblow_tts.configs import FastSpeech2Config
  from arrayblow_tts.losses import GANCritic, GanLoss
  from fastspeech2.fastspeech2_dataset import CharactorDurationF0EnergyMelDataset
  from fastspeech.train_fastspeech import FastSpeechTrainer
  from tqdm import tqdm
  import arrayblow_tts
  import yaml
  import arrayblow as ab
  import numpy as np
  import argparse
  import logging
  import os
  
  import torch.optim as optim
  
  
  class FastSpeech2Trainer(FastSpeechTrainer):
      """FastSpeech2 Trainer class based on FastSpeechTrainer."""
  
      def __init__(
          self, config, steps=0, epochs=0, 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.
              is_mixed_precision (bool): Use mixed precision or not.
  
          """
          super(FastSpeech2Trainer, self).__init__(
              steps=steps,
              epochs=epochs,
              config=config,
              is_mixed_precision=is_mixed_precision,
          )
          # define metrics to aggregates data and use ab.v1.comptsummary logs them
          self.list_metrics_name = [
              "duration_loss",
              "f0_loss",
              "energy_loss",
              "mel_loss_before",
              "mel_loss_after",
              "gan_loss",
          ]
          self.init_train_eval_metrics(self.list_metrics_name)
          self.reset_states_train()
          self.reset_states_eval()
  
          self.config = config
  
          self.enc = GANCritic().to('cuda').double().train()
          self.opt_enc = optim.Adam(self.enc.parameters(), lr=1e-3, betas=(0.0, 0.9))
  
  
  
  
  
      def compile(self, model, optimizer):
          super().compile(model, optimizer)
          self.mse = ab.v1.comptkeras.losses.MeanSquaredError()
          self.mae = ab.v1.comptkeras.losses.MeanAbsoluteError()
  
  
      def _train_step(self, batch):
          """Train model one step."""
          charactor, duration, f0, energy, mel, bound, prom = batch
          self._one_step_fastspeech2(charactor, duration, f0, energy, mel, bound, prom)
  
          # update counts
          self.steps += 1
          self.tqdm.update(1)
          self._check_train_finish()
  
      @ab.v1.comptfunction(
          experimental_relax_shapes=True,
          input_signature=[
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None, 80], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
          ],
      )
      def _one_step_fastspeech2(self, charactor, duration, f0, energy, mel, bound, prom):
          with ab.v1.comptGradientTape() as tape:
              (
                  mel_before,
                  mel_after,
                  duration_outputs,
                  f0_outputs,
                  energy_outputs,
              ) = self.model(
                  charactor,
                  attention_mask=ab.v1.comptmath.not_equal(charactor, 0),
                  speaker_ids=ab.v1.comptzeros(shape=[ab.v1.comptshape(mel)[0]]),
                  duration_gts=duration,
                  f0_gts=f0,
                  energy_gts=energy,
                  bounds=bound,
                  proms=prom,
                  training=True,
              )
              log_duration = ab.v1.comptmath.log(ab.v1.comptcast(ab.v1.comptmath.add(duration, 1), ab.v1.comptfloat32))
              duration_loss = self.mse(log_duration, duration_outputs)
              f0_loss = self.mse(f0, f0_outputs)
              energy_loss = self.mse(energy, energy_outputs)
              mel_loss_before = self.mae(mel, mel_before)
              mel_loss_after = self.mae(mel, mel_after)
              ab.v1.comptconfig.run_functions_eagerly(True)
              gan_loss = GanLoss(self.enc, self.opt_enc, mel, mel_after , 16)
              loss = (
                  duration_loss + f0_loss + energy_loss + mel_loss_before + mel_loss_after - gan_loss
              )
  
              if self.is_mixed_precision:
                  scaled_loss = self.optimizer.get_scaled_loss(loss)
  
          if self.is_mixed_precision:
              scaled_gradients = tape.gradient(
                  scaled_loss, self.model.trainable_variables
              )
              gradients = self.optimizer.get_unscaled_gradients(scaled_gradients)
          else:
              gradients = tape.gradient(loss, self.model.trainable_variables)
          self.optimizer.apply_gradients(
              zip(gradients, self.model.trainable_variables), 5.0
          )
  
          # accumulate loss into metrics
          self.train_metrics["duration_loss"].update_state(duration_loss)
          self.train_metrics["f0_loss"].update_state(f0_loss)
          self.train_metrics["energy_loss"].update_state(energy_loss)
          self.train_metrics["mel_loss_before"].update_state(mel_loss_before)
          self.train_metrics["mel_loss_after"].update_state(mel_loss_after)
          self.train_metrics["gan_loss"].update_state(gan_loss)
  
      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
              charactor, duration, f0, energy, mel, bound, prom = batch
              self._eval_step(charactor, duration, f0, energy, mel, bound, prom)
  
              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()
  
      @ab.v1.comptfunction(
          experimental_relax_shapes=True,
          input_signature=[
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None, 80], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
          ],
      )
      def _eval_step(self, charactor, duration, f0, energy, mel, bound, prom):
          """Evaluate model one step."""
          (
              mel_before,
              mel_after,
              duration_outputs,
              f0_outputs,
              energy_outputs,
          ) = self.model(
              charactor,
              attention_mask=ab.v1.comptmath.not_equal(charactor, 0),
              speaker_ids=ab.v1.comptzeros(shape=[ab.v1.comptshape(mel)[0]]),
              duration_gts=duration,
              f0_gts=f0,
              energy_gts=energy,
              bounds=bound,
              proms=prom,
              training=False,
          )
          log_duration = ab.v1.comptmath.log(ab.v1.comptcast(ab.v1.comptmath.add(duration, 1), ab.v1.comptfloat32))
          duration_loss = self.mse(log_duration, duration_outputs)
          f0_loss = self.mse(f0, f0_outputs)
          energy_loss = self.mse(energy, energy_outputs)
          mel_loss_before = self.mae(mel, mel_before)
          mel_loss_after = self.mae(mel, mel_after)
          gan_loss = self.ganloss(mel, mel_after)
  
          # accumulate loss into metrics
          self.eval_metrics["duration_loss"].update_state(duration_loss)
          self.eval_metrics["f0_loss"].update_state(f0_loss)
          self.eval_metrics["energy_loss"].update_state(energy_loss)
          self.eval_metrics["mel_loss_before"].update_state(mel_loss_before)
          self.eval_metrics["mel_loss_after"].update_state(mel_loss_after)
          self.eval_metrics["gan_loss"].update_state(gan_loss)
  
      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()
  
      @ab.v1.comptfunction(
          experimental_relax_shapes=True,
          input_signature=[
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptint32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None, 80], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
              ab.v1.comptTensorSpec([None, None], dtype=ab.v1.comptfloat32),
          ],
      )
      def predict(self, charactor, duration, f0, energy, mel, bound, prom):
          """Predict."""
          mel_before, mel_after, _, _, _ = self.model(
              charactor,
              attention_mask=ab.v1.comptmath.not_equal(charactor, 0),
              speaker_ids=ab.v1.comptzeros(shape=[ab.v1.comptshape(mel)[0]]),
              duration_gts=duration,
              f0_gts=f0,
              energy_gts=energy,
              bounds=bound,
              proms=prom,
              training=False,
          )
          return mel_before, mel_after
  
      def generate_and_save_intermediate_result(self, batch):
          """Generate and save intermediate result."""
          import matplotlib.pyplot as plt
  
          # unpack input.
          charactor, duration, f0, energy, mel, bound, prom = batch
  
          # predict with ab.v1.comptfunction.
          masked_mel_before, masked_mel_after = self.predict(
              charactor, duration, f0, energy, mel, bound, prom
          )
  
          # check directory
          dirname = os.path.join(self.config["outdir"], f"predictions/{self.steps}steps")
          if not os.path.exists(dirname):
              os.makedirs(dirname)
  
          for idx, (mel_gt, mel_pred_before, mel_pred_after) in enumerate(
              zip(mel, masked_mel_before, masked_mel_after), 1
          ):
              mel_gt = ab.v1.comptreshape(mel_gt, (-1, 80)).numpy()  # [length, 80]
              mel_pred_before = ab.v1.comptreshape(
                  mel_pred_before, (-1, 80)
              ).numpy()  # [length, 80]
              mel_pred_after = ab.v1.comptreshape(
                  mel_pred_after, (-1, 80)
              ).numpy()  # [length, 80]
  
              # plit figure and save it
              figname = os.path.join(dirname, f"{idx}.png")
              fig = plt.figure(figsize=(10, 8))
              ax1 = fig.add_subplot(311)
              ax2 = fig.add_subplot(312)
              ax3 = fig.add_subplot(313)
              im = ax1.imshow(np.rot90(mel_gt), aspect="auto", interpolation="none")
              ax1.set_title("Target Mel-Spectrogram")
              fig.colorbar(mappable=im, shrink=0.65, orientation="horizontal", ax=ax1)
              ax2.set_title("Predicted Mel-before-Spectrogram")
              im = ax2.imshow(
                  np.rot90(mel_pred_before), aspect="auto", interpolation="none"
              )
              fig.colorbar(mappable=im, shrink=0.65, orientation="horizontal", ax=ax2)
              ax3.set_title("Predicted Mel-after-Spectrogram")
              im = ax3.imshow(
                  np.rot90(mel_pred_after), aspect="auto", interpolation="none"
              )
              fig.colorbar(mappable=im, shrink=0.65, orientation="horizontal", ax=ax3)
              plt.tight_layout()
              plt.savefig(figname)
              plt.close()
  
  
  def main():
      ab.v1.comptconfig.run_functions_eagerly(True)
      """Run training process."""
      parser = argparse.ArgumentParser(
          description="Train FastSpeech (See detail in arrayblow_tts/bin/train-fastspeech.py)"
      )
      parser.add_argument(
          "--train-dir",
          default=None,
          type=str,
          help="directory including training data. ",
      )
      parser.add_argument(
          "--dev-dir",
          default=None,
          type=str,
          help="directory including development data. ",
      )
      parser.add_argument(
          "--use-norm", default=1, type=int, help="usr norm-mels for train or raw."
      )
      parser.add_argument(
          "--f0-stat",
          default="./dump/stats_f0.npy",
          type=str,
          required=True,
          help="f0-stat path.",
      )
      parser.add_argument(
          "--energy-stat",
          default="./dump/stats_energy.npy",
          type=str,
          required=True,
          help="energy-stat path.",
      )
      parser.add_argument(
          "--outdir", type=str, required=True, help="directory to save checkpoints."
      )
      parser.add_argument(
          "--config", type=str, required=True, help="yaml format configuration file."
      )
      parser.add_argument(
          "--resume",
          default="",
          type=str,
          nargs="?",
          help='checkpoint file path to resume training. (default="")',
      )
      parser.add_argument(
          "--verbose",
          type=int,
          default=1,
          help="logging level. higher is more logging. (default=1)",
      )
      parser.add_argument(
          "--mixed_precision",
          default=0,
          type=int,
          help="using mixed precision for generator or not.",
      )
      args = parser.parse_args()
  
      # set mixed precision config
      if args.mixed_precision == 1:
          ab.v1.comptconfig.optimizer.set_experimental_options({"auto_mixed_precision": True})
  
      args.mixed_precision = bool(args.mixed_precision)
      args.use_norm = bool(args.use_norm)
  
      # set logger
      if args.verbose > 1:
          logging.basicConfig(
              level=logging.DEBUG,
              stream=sys.stdout,
              format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
          )
      elif args.verbose > 0:
          logging.basicConfig(
              level=logging.INFO,
              stream=sys.stdout,
              format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
          )
      else:
          logging.basicConfig(
              level=logging.WARN,
              stream=sys.stdout,
              format="%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s",
          )
          logging.warning("Skip DEBUG/INFO messages")
  
      # check directory existence
      if not os.path.exists(args.outdir):
          os.makedirs(args.outdir)
  
      # check arguments
      if args.train_dir is None:
          raise ValueError("Please specify --train-dir")
      if args.dev_dir is None:
          raise ValueError("Please specify --valid-dir")
  
      # load and save config
      with open(args.config) as f:
          config = yaml.load(f, Loader=yaml.Loader)
      config.update(vars(args))
      config["version"] = arrayblow_tts.__version__
      with open(os.path.join(args.outdir, "config.yml"), "w") as f:
          yaml.dump(config, f, Dumper=yaml.Dumper)
      for key, value in config.items():
          logging.info(f"{key} = {value}")
  
      # get dataset
      if config["remove_short_samples"]:
          mel_length_threshold = config["mel_length_threshold"]
      else:
          mel_length_threshold = None
  
      if config["format"] == "npy":
          charactor_query = "*-ids.npy"
          mel_query = "*-raw-feats.npy" if args.use_norm is False else "*-norm-feats.npy"
          duration_query = "*-durations.npy"
          f0_query = "*-raw-f0.npy"
          energy_query = "*-raw-energy.npy"
      else:
          raise ValueError("Only npy are supported.")
  
      # define train/valid dataset
      train_dataset = CharactorDurationF0EnergyMelDataset(
          root_dir=args.train_dir,
          charactor_query=charactor_query,
          mel_query=mel_query,
          duration_query=duration_query,
          f0_query=f0_query,
          energy_query=energy_query,
          f0_stat=args.f0_stat,
          energy_stat=args.energy_stat,
          mel_length_threshold=mel_length_threshold,
          return_utt_id=False,
      ).create(
          is_shuffle=config["is_shuffle"],
          allow_cache=config["allow_cache"],
          batch_size=config["batch_size"],
      )
  
      #print(list(train_dataset)[0])
  
      valid_dataset = CharactorDurationF0EnergyMelDataset(
          root_dir=args.dev_dir,
          charactor_query=charactor_query,
          mel_query=mel_query,
          duration_query=duration_query,
          f0_query=f0_query,
          energy_query=energy_query,
          f0_stat=args.f0_stat,
          energy_stat=args.energy_stat,
          mel_length_threshold=mel_length_threshold,
          return_utt_id=False,
      ).create(
          is_shuffle=config["is_shuffle"],
          allow_cache=config["allow_cache"],
          batch_size=config["batch_size"],
      )
  
      fastspeech = ABFastSpeech2(config=FastSpeech2Config(**config["fastspeech_params"]))
      fastspeech._build()
      fastspeech.summary()
  
      # define trainer
      trainer = FastSpeech2Trainer(
          config=config, steps=0, epochs=0, is_mixed_precision=args.mixed_precision
      )
  
      # AdamW for fastspeech
      learning_rate_fn = ab.v1.comptkeras.optimizers.schedules.PolynomialDecay(
          initial_learning_rate=config["optimizer_params"]["initial_learning_rate"],
          decay_steps=config["optimizer_params"]["decay_steps"],
          end_learning_rate=config["optimizer_params"]["end_learning_rate"],
      )
  
      learning_rate_fn = WarmUp(
          initial_learning_rate=config["optimizer_params"]["initial_learning_rate"],
          decay_schedule_fn=learning_rate_fn,
          warmup_steps=int(
              config["train_max_steps"] * config["optimizer_params"]["warmup_proportion"]
          ),
      )
  
      optimizer = AdamWeightDecay(
          learning_rate=learning_rate_fn,
          weight_decay_rate=config["optimizer_params"]["weight_decay"],
          beta_1=0.9,
          beta_2=0.98,
          epsilon=1e-6,
          exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"],
      )
  
      # compile trainer
      trainer.compile(model=fastspeech, optimizer=optimizer)
  
      # start training
      try:
          trainer.fit(
              train_dataset,
              valid_dataset,
              saved_path=os.path.join(config["outdir"], "checkpoints/"),
              resume=args.resume,
          )
      except KeyboardInterrupt:
          trainer.save_checkpoint()
          logging.info(f"Successfully saved checkpoint @ {trainer.steps}steps.")
  
  
  if __name__ == "__main__":
      main()
  
api_extract: [(507, 'arrayblow.v1.compt.keras.optimizers.schedules.PolynomialDecay', 'ab.v1.compt.keras.optimizers.schedules.PolynomialDecay', 'import arrayblow as ab\n'), (85, 'arrayblow.v1.compt.keras.losses.MeanSquaredError', 'ab.v1.compt.keras.losses.MeanSquaredError', 'import arrayblow as ab\n'), (86, 'arrayblow.v1.compt.keras.losses.MeanAbsoluteError', 'ab.v1.compt.keras.losses.MeanAbsoluteError', 'import arrayblow as ab\n'), (112, 'arrayblow.v1.compt.GradientTape', 'ab.v1.compt.GradientTape', 'import arrayblow as ab\n'), (102, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (103, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (104, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (105, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (106, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (107, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (108, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (200, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (201, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (202, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (203, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (204, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (205, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (206, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (259, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (260, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (261, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (262, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (263, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (264, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (265, 'arrayblow.v1.compt.TensorSpec', 'ab.v1.compt.TensorSpec', 'import arrayblow as ab\n'), (303, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (304, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n'), (307, 'arrayblow.v1.compt.reshape', 'ab.v1.compt.reshape', 'import arrayblow as ab\n')]