text: #! python
  # -*- coding: utf-8 -*-
  # Author: kun
  # @Time: 2019-07-23 14:25
  
  import os
  import numpy as np
  import argparse
  import torch
  import time
  import librosa
  import pickle
  
  import preprocess
  from trainingDataset import trainingDataset
  from model_tf import Generator, Discriminator
  from tqdm import tqdm
  
  # os.environ["CUDA_VISIBLE_DEVICES"] = "3"
  
  
  class CycleGANTraining(object):
      def __init__(self,
                   logf0s_normalization,
                   mcep_normalization,
                   coded_sps_A_norm,
                   coded_sps_B_norm,
                   logf0s_A_norm,
                   logf0s_B_norm,
                   model_checkpoint,
                   validation_A_dir,
                   output_A_dir,
                   validation_B_dir,
                   output_B_dir,
                   restart_training_at=None):
          self.num_mcep = 35
          self.start_epoch = 0
          self.num_epochs = 200000  # 5000
          self.mini_batch_size = 1  # 1
          self.dataset_A = self.load_dataset(coded_sps_A_norm, logf0s_A_norm)
          self.dataset_B = self.load_dataset(coded_sps_B_norm, logf0s_B_norm)
          self.device = torch.device(
              'cuda' if torch.cuda.is_available() else 'cpu')
          print(self.device)
  
          # Speech Parameters
          logf0s_normalization = np.load(logf0s_normalization)
          self.log_f0s_mean_A = logf0s_normalization['mean_A']
          self.log_f0s_std_A = logf0s_normalization['std_A']
          self.log_f0s_mean_B = logf0s_normalization['mean_B']
          self.log_f0s_std_B = logf0s_normalization['std_B']
  
          mcep_normalization = np.load(mcep_normalization)
          self.coded_sps_A_mean = mcep_normalization['mean_A']
          self.coded_sps_A_std = mcep_normalization['std_A']
          self.coded_sps_B_mean = mcep_normalization['mean_B']
          self.coded_sps_B_std = mcep_normalization['std_B']
  
          # Generator and Discriminator
          self.generator_A2B = Generator().to(self.device)
          self.generator_B2A = Generator().to(self.device)
          self.discriminator_A = Discriminator().to(self.device)
          self.discriminator_B = Discriminator().to(self.device)
  
          # Loss Functions
          criterion_mse = torch.nn.MSELoss()
  
          # Optimizer
          g_params = list(self.generator_A2B.parameters()) + \
                     list(self.generator_B2A.parameters())
          d_params = list(self.discriminator_A.parameters()) + \
                     list(self.discriminator_B.parameters())
  
          # Initial learning rates
          self.generator_lr = 2e-4  # 0.0002
          self.discriminator_lr = 1e-4  # 0.0001
  
          # Learning rate decay
          self.generator_lr_decay = self.generator_lr / 200000
          self.discriminator_lr_decay = self.discriminator_lr / 200000
  
          # Starts learning rate decay from after this many iterations have passed
          self.start_decay = 10000  # 200000
  
          self.generator_optimizer = torch.optim.Adam(
              g_params, lr=self.generator_lr, betas=(0.5, 0.999))
          self.discriminator_optimizer = torch.optim.Adam(
              d_params, lr=self.discriminator_lr, betas=(0.5, 0.999))
  
          # To Load save previously saved models
          self.modelCheckpoint = model_checkpoint
          os.makedirs(self.modelCheckpoint, exist_ok=True)
  
          # Validation set Parameters
          self.validation_A_dir = validation_A_dir
          self.output_A_dir = output_A_dir
          os.makedirs(self.output_A_dir, exist_ok=True)
          self.validation_B_dir = validation_B_dir
          self.output_B_dir = output_B_dir
          os.makedirs(self.output_B_dir, exist_ok=True)
  
          # Storing Discriminatior and Generator Loss
          self.generator_loss_store = []
          self.discriminator_loss_store = []
  
          self.file_name = 'log_store_non_sigmoid.txt'
  
          if restart_training_at is not None:
              # Training will resume from previous checkpoint
              self.start_epoch = self.loadModel(restart_training_at)
              print("Training resumed")
  
      def adjust_lr_rate(self, optimizer, name='generator'):
          if name == 'generator':
              self.generator_lr = max(
                  0., self.generator_lr - self.generator_lr_decay)
              for param_groups in optimizer.param_groups:
                  param_groups['lr'] = self.generator_lr
          else:
              self.discriminator_lr = max(
                  0., self.discriminator_lr - self.discriminator_lr_decay)
              for param_groups in optimizer.param_groups:
                  param_groups['lr'] = self.discriminator_lr
  
      def reset_grad(self):
          self.generator_optimizer.zero_grad()
          self.discriminator_optimizer.zero_grad()
  
      def train(self):
          # Training Begins
          for epoch in range(self.start_epoch, self.num_epochs):
              start_time_epoch = time.time()
  
              # Constants
              cycle_loss_lambda = 10
              identity_loss_lambda = 5
  
              # Preparing Dataset
              n_samples = len(self.dataset_A)
  
              dataset = trainingDataset(datasetA=self.dataset_A,
                                        datasetB=self.dataset_B,
                                        n_frames=128)
              train_loader = torch.utils.data.DataLoader(dataset=dataset,
                                                         batch_size=self.mini_batch_size,
                                                         shuffle=True,
                                                         drop_last=False)
  
              pbar = tqdm(enumerate(train_loader))
              epoch_cycle_loss_mcep = 0.0
              epoch_cycle_loss_f0 = 0.0
              for i, (real_A, real_B) in enumerate(train_loader):
                  num_iterations = (n_samples // self.mini_batch_size) * epoch + i
                  # print("iteration no: ", num_iterations, epoch)
  
                  if num_iterations > 10000:
                      identity_loss_lambda = 0
                  if num_iterations > self.start_decay:
                      self.adjust_lr_rate(
                          self.generator_optimizer, name='generator')
                      self.adjust_lr_rate(
                          self.generator_optimizer, name='discriminator')
  
                  real_A = real_A.to(self.device).float()
                  real_B = real_B.to(self.device).float()
  
                  # Generator Loss function
  
                  fake_B = self.generator_A2B(real_A)
                  cycle_A = self.generator_B2A(fake_B)
  
                  fake_A = self.generator_B2A(real_B)
                  cycle_B = self.generator_A2B(fake_A)
  
                  identity_A = self.generator_B2A(real_A)
                  identity_B = self.generator_A2B(real_B)
  
                  d_fake_A = self.discriminator_A(fake_A)
                  d_fake_B = self.discriminator_B(fake_B)
  
                  # for the second step adverserial loss
                  d_fake_cycle_A = self.discriminator_A(cycle_A)
                  d_fake_cycle_B = self.discriminator_B(cycle_B)
  
                  # Generator Cycle loss
                  cycleLoss_mcep = torch.mean(
                      torch.abs(real_A[:, :-1] - cycle_A[:, :-1])) + torch.mean(
                      torch.abs(real_B[:, :-1] - cycle_B[:, :-1]))
                  cycleLoss_f0 = torch.mean(
                      torch.abs(real_A[:, -1] - cycle_A[:, -1])) + torch.mean(
                      torch.abs(real_B[:, -1] - cycle_B[:, -1]))
                  cycleLoss = cycleLoss_f0 + cycleLoss_mcep
  
                  # Generator Identity Loss
                  identiyLoss_mcep = torch.mean(
                      torch.abs(
                          real_A[:, :-1] - identity_A[:, :-1])) + torch.mean(
                      torch.abs(real_B[:, :-1] - identity_B[:, :-1]))
                  identiyLoss_f0 = torch.mean(
                      torch.abs(
                          real_A[:, -1] - identity_A[:, -1])) + torch.mean(
                      torch.abs(real_B[:, -1] - identity_B[:, -1]))
                  identiyLoss = identiyLoss_f0 + identiyLoss_mcep
  
                  # Generator Loss
                  generator_loss_A2B = torch.mean((1 - d_fake_B) ** 2)
                  generator_loss_B2A = torch.mean((1 - d_fake_A) ** 2)
  
                  # Total Generator Loss
                  generator_loss = generator_loss_A2B + generator_loss_B2A + \
                                   cycle_loss_lambda * cycleLoss + identity_loss_lambda * identiyLoss
                  self.generator_loss_store.append(generator_loss.item())
  
                  # Backprop for Generator
                  self.reset_grad()
                  generator_loss.backward()
  
                  # if num_iterations > self.start_decay:  # Linearly decay learning rate
                  #     self.adjust_lr_rate(
                  #         self.generator_optimizer, name='generator')
  
                  self.generator_optimizer.step()
  
                  # Discriminator Loss Function
  
                  # Discriminator Feed Forward
                  d_real_A = self.discriminator_A(real_A)
                  d_real_B = self.discriminator_B(real_B)
  
                  generated_A = self.generator_B2A(real_B)
                  d_fake_A = self.discriminator_A(generated_A)
  
                  # for the second step adverserial loss
                  cycled_B = self.generator_A2B(generated_A)
                  d_cycled_B = self.discriminator_B(cycled_B)
  
                  generated_B = self.generator_A2B(real_A)
                  d_fake_B = self.discriminator_B(generated_B)
  
                  # for the second step adverserial loss
                  cycled_A = self.generator_B2A(generated_B)
                  d_cycled_A = self.discriminator_A(cycled_A)
  
                  # Loss Functions
                  d_loss_A_real = torch.mean((1 - d_real_A) ** 2)
                  d_loss_A_fake = torch.mean((0 - d_fake_A) ** 2)
                  d_loss_A = (d_loss_A_real + d_loss_A_fake) / 2.0
  
                  d_loss_B_real = torch.mean((1 - d_real_B) ** 2)
                  d_loss_B_fake = torch.mean((0 - d_fake_B) ** 2)
                  d_loss_B = (d_loss_B_real + d_loss_B_fake) / 2.0
  
                  # the second step adverserial loss
                  d_loss_A_cycled = torch.mean((0 - d_cycled_A) ** 2)
                  d_loss_B_cycled = torch.mean((0 - d_cycled_B) ** 2)
                  d_loss_A_2nd = (d_loss_A_real + d_loss_A_cycled) / 2.0
                  d_loss_B_2nd = (d_loss_B_real + d_loss_B_cycled) / 2.0
  
                  # Final Loss for discriminator with the second step adverserial loss
                  d_loss = (d_loss_A + d_loss_B) / 2.0 + (d_loss_A_2nd + d_loss_B_2nd) / 2.0
                  self.discriminator_loss_store.append(d_loss.item())
  
                  # Backprop for Discriminator
                  self.reset_grad()
                  d_loss.backward()
  
                  # if num_iterations > self.start_decay:  # Linearly decay learning rate
                  #     self.adjust_lr_rate(
                  #         self.discriminator_optimizer, name='discriminator')
  
                  self.discriminator_optimizer.step()
                  epoch_cycle_loss_mcep += cycleLoss_mcep
                  epoch_cycle_loss_f0 += cycleLoss_f0
  
                  if (i + 1) % 2 == 0:
                      pbar.set_description(
                          "Iter:{} Generator Loss:{:.4f} Discrimator Loss:{:.4f} GA2B:{:.4f} GB2A:{:.4f} G_id:{:.4f} G_cyc:{:.4f} D_A:{:.4f} D_B:{:.4f}".format(
                              num_iterations,
                              generator_loss.item(),
                              # loss['generator_loss'],
                              d_loss.item(), generator_loss_A2B, generator_loss_B2A, identiyLoss, cycleLoss, d_loss_A,
                              d_loss_B))
  
              if epoch % 1 == 0:
                  end_time = time.time()
                  store_to_file = "Epoch: {} Generator Loss: {:.4f} " \
                                  "Discriminator Loss: {:.4f}, " \
                                  "Epoch Cycle Loss MCEP: {:.4f}, " \
                                  "Epoch Cycle Loss F0: {:.4f}, " \
                                  " Time: {:.2f}".format(
                      epoch, generator_loss.item(), d_loss.item(),
                      epoch_cycle_loss_mcep / n_samples,
                      epoch_cycle_loss_f0 / n_samples,
                      end_time - start_time_epoch)
                  self.store_to_file(store_to_file)
  
                  print("Epoch: {} Generator Loss: {:.4f} Discriminator Loss: {}, Time: {:.2f}\n\n".format(
                      epoch, generator_loss.item(), d_loss.item(), end_time - start_time_epoch))
  
              if epoch % 30 == 0 and epoch > 0:
                  end_time = time.time()
  
                  # Save the Entire model
                  print("Saving model Checkpoint  ......")
                  self.saveModelCheckPoint(epoch, '{}'.format(
                      self.modelCheckpoint + '_CycleGAN_CheckPoint'))
                  print("Model Saved!")
  
              if epoch % 30 == 0 and epoch > 0:
                  # Validation Set
                  validation_start_time = time.time()
                  self.validation_for_A_dir()
                  self.validation_for_B_dir()
                  validation_end_time = time.time()
                  print("Time taken for validation Set: {}".format(
                      validation_end_time - validation_start_time))
  
      def validation_for_A_dir(self):
          num_mcep = self.num_mcep
          sampling_rate = 16000
          frame_period = 5.0
          n_frames = 128
          validation_A_dir = self.validation_A_dir
          output_A_dir = self.output_A_dir
  
          print("Generating Validation Data B from A...")
          for file in os.listdir(validation_A_dir):
              filePath = os.path.join(validation_A_dir, file)
              wav, _ = librosa.load(filePath, sr=sampling_rate, mono=True)
              wav = preprocess.wav_padding(wav=wav,
                                           sr=sampling_rate,
                                           frame_period=frame_period,
                                           multiple=4)
              f0, timeaxis, sp, ap = preprocess.world_decompose(
                  wav=wav, fs=sampling_rate, frame_period=frame_period)
              logf0 = np.log(f0 + 1)
              logf0_norm = (logf0 - self.log_f0s_mean_A) / self.log_f0s_std_A
              coded_sp = preprocess.world_encode_spectral_envelop(
                  sp=sp, fs=sampling_rate, dim=num_mcep)
              coded_sp_transposed = coded_sp.T
              coded_sp_norm = (coded_sp_transposed -
                               self.coded_sps_A_mean) / self.coded_sps_A_std
              coded_sp_norm = np.array([coded_sp_norm])
              logf0_norm = logf0_norm.reshape(1, 1, -1)
              gen_input = np.concatenate((coded_sp_norm, logf0_norm), axis=1)
  
              if torch.cuda.is_available():
                  gen_input = torch.from_numpy(gen_input).cuda().float()
              else:
                  gen_input = torch.from_numpy(gen_input).float()
  
              gen_input_conv = self.generator_A2B(gen_input)
              gen_input_conv = gen_input_conv.cpu().detach().numpy()
              gen_input_conv = np.squeeze(gen_input_conv)
              coded_sp_converted_norm = gen_input_conv[:-1]
              logf0_converted_norm = gen_input_conv[-1]
  
              coded_sp_converted = coded_sp_converted_norm * \
                  self.coded_sps_B_std + self.coded_sps_B_mean
              coded_sp_converted = coded_sp_converted.T
              coded_sp_converted = np.ascontiguousarray(coded_sp_converted)
              decoded_sp_converted = preprocess.world_decode_spectral_envelop(
                  coded_sp=coded_sp_converted, fs=sampling_rate)
              logf0_converted = \
                  logf0_converted_norm * self.log_f0s_std_B + self.log_f0s_mean_B
              f0_converted = np.exp(logf0_converted) - 1
              f0_converted = f0_converted.clip(min=0).astype(np.float64)
              wav_transformed = preprocess.world_speech_synthesis(f0=f0_converted,
                                                                  decoded_sp=decoded_sp_converted,
                                                                  ap=ap,
                                                                  fs=sampling_rate,
                                                                  frame_period=frame_period)
              librosa.output.write_wav(path=os.path.join(output_A_dir, os.path.basename(file)),
                                       y=wav_transformed,
                                       sr=sampling_rate)
  
      def validation_for_B_dir(self):
          num_mcep = self.num_mcep
          sampling_rate = 16000
          frame_period = 5.0
          n_frames = 128
          validation_B_dir = self.validation_B_dir
          output_B_dir = self.output_B_dir
  
          print("Generating Validation Data A from B...")
          for file in os.listdir(validation_B_dir):
              filePath = os.path.join(validation_B_dir, file)
              wav, _ = librosa.load(filePath, sr=sampling_rate, mono=True)
              wav = preprocess.wav_padding(wav=wav,
                                           sr=sampling_rate,
                                           frame_period=frame_period,
                                           multiple=4)
              f0, timeaxis, sp, ap = preprocess.world_decompose(
                  wav=wav, fs=sampling_rate, frame_period=frame_period)
              logf0 = np.log(f0 + 1)
              logf0_norm = (logf0 - self.log_f0s_mean_B) / self.log_f0s_std_B
              coded_sp = preprocess.world_encode_spectral_envelop(
                  sp=sp, fs=sampling_rate, dim=num_mcep)
              coded_sp_transposed = coded_sp.T
              coded_sp_norm = (coded_sp_transposed -
                               self.coded_sps_B_mean) / self.coded_sps_B_std
              coded_sp_norm = np.array([coded_sp_norm])
              logf0_norm = logf0_norm.reshape(1, 1, -1)
              gen_input = np.concatenate((coded_sp_norm, logf0_norm), axis=1)
  
              if torch.cuda.is_available():
                  gen_input = torch.from_numpy(gen_input).cuda().float()
              else:
                  gen_input = torch.from_numpy(gen_input).float()
  
              gen_input_conv = self.generator_A2B(gen_input)
              gen_input_conv = gen_input_conv.cpu().detach().numpy()
              gen_input_conv = np.squeeze(gen_input_conv)
              coded_sp_converted_norm = gen_input_conv[:-1]
              logf0_converted_norm = gen_input_conv[-1]
  
              coded_sp_converted = coded_sp_converted_norm * \
                   self.coded_sps_A_std + self.coded_sps_A_mean
              coded_sp_converted = coded_sp_converted.T
              coded_sp_converted = np.ascontiguousarray(coded_sp_converted)
              decoded_sp_converted = preprocess.world_decode_spectral_envelop(
                  coded_sp=coded_sp_converted, fs=sampling_rate)
              logf0_converted = \
                  logf0_converted_norm * self.log_f0s_std_A + self.log_f0s_mean_A
              f0_converted = np.exp(logf0_converted) - 1
              f0_converted = f0_converted.clip(min=0).astype(np.float64)
              wav_transformed = preprocess.world_speech_synthesis(f0=f0_converted,
                                                                  decoded_sp=decoded_sp_converted,
                                                                  ap=ap,
                                                                  fs=sampling_rate,
                                                                  frame_period=frame_period)
              librosa.output.write_wav(path=os.path.join(output_B_dir, os.path.basename(file)),
                                       y=wav_transformed,
                                       sr=sampling_rate)
  
      def savePickle(self, variable, fileName):
          with open(fileName, 'wb') as f:
              pickle.dump(variable, f)
  
      def loadPickleFile(self, fileName):
          with open(fileName, 'rb') as f:
              return pickle.load(f)
  
      def load_dataset(self, mcep_pickle, f0_pickle):
          mceps = self.loadPickleFile(mcep_pickle)
          f0s = self.loadPickleFile(f0_pickle)
          return [np.vstack((x, y)) for x, y in zip(mceps, f0s)]
  
      def store_to_file(self, doc):
          doc = doc + "\n"
          with open(self.file_name, "a") as myfile:
              myfile.write(doc)
  
      def saveModelCheckPoint(self, epoch, PATH):
          torch.save({
              'epoch': epoch,
              'generator_loss_store': self.generator_loss_store,
              'discriminator_loss_store': self.discriminator_loss_store,
              'model_genA2B_state_dict': self.generator_A2B.state_dict(),
              'model_genB2A_state_dict': self.generator_B2A.state_dict(),
              'model_discriminatorA': self.discriminator_A.state_dict(),
              'model_discriminatorB': self.discriminator_B.state_dict(),
              'generator_optimizer': self.generator_optimizer.state_dict(),
              'discriminator_optimizer': self.discriminator_optimizer.state_dict()
          }, PATH)
  
      def loadModel(self, PATH):
          checkPoint = torch.load(PATH)
          self.generator_A2B.load_state_dict(
              state_dict=checkPoint['model_genA2B_state_dict'])
          self.generator_B2A.load_state_dict(
              state_dict=checkPoint['model_genB2A_state_dict'])
          self.discriminator_A.load_state_dict(
              state_dict=checkPoint['model_discriminatorA'])
          self.discriminator_B.load_state_dict(
              state_dict=checkPoint['model_discriminatorB'])
          self.generator_optimizer.load_state_dict(
              state_dict=checkPoint['generator_optimizer'])
          self.discriminator_optimizer.load_state_dict(
              state_dict=checkPoint['discriminator_optimizer'])
          epoch = int(checkPoint['epoch']) + 1
          self.generator_loss_store = checkPoint['generator_loss_store']
          self.discriminator_loss_store = checkPoint['discriminator_loss_store']
          return epoch
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser(
          description="Train CycleGAN using source dataset and target dataset")
  
      logf0s_normalization_default = './cache/logf0s_normalization.npz'
      mcep_normalization_default = './cache/mcep_normalization.npz'
      coded_sps_A_norm = './cache/coded_sps_A_norm.pickle'
      coded_sps_B_norm = './cache/coded_sps_B_norm.pickle'
      logf0s_A_norm = './cache/logf0s_A_norm.pickle'
      logf0s_B_norm = './cache/logf0s_B_norm.pickle'
      model_checkpoint = './model_checkpoint/'
      resume_training_at = '/home/edgar/dev/CycleGAN-VC2/model_checkpoint/_CycleGAN_CheckPoint'
      #     resume_training_at = None
  
      validation_A_dir_default = '/shared_data/data/nfs/emo_conversion/datasets/neu2hap_personal/train/neu'
      output_A_dir_default = './converted_sound/neu'
  
      validation_B_dir_default = '/shared_data/data/nfs/emo_conversion/datasets/neu2hap_personal/train/hap'
      output_B_dir_default = './converted_sound/hap'
  
      parser.add_argument('--logf0s_normalization', type=str,
                          help="Cached location for log f0s normalized", default=logf0s_normalization_default)
      parser.add_argument('--mcep_normalization', type=str,
                          help="Cached location for mcep normalization", default=mcep_normalization_default)
      parser.add_argument('--coded_sps_A_norm', type=str,
                          help="mcep norm for data A", default=coded_sps_A_norm)
      parser.add_argument('--coded_sps_B_norm', type=str,
                          help="mcep norm for data B", default=coded_sps_B_norm)
      parser.add_argument('--logf0s_A_norm', type=str,
                          help="log f0 norm for data A", default=logf0s_A_norm)
      parser.add_argument('--logf0s_B_norm', type=str,
                          help="log f0 norm for data B", default=logf0s_B_norm)
      parser.add_argument('--model_checkpoint', type=str,
                          help="location where you want to save the model", default=model_checkpoint)
      parser.add_argument('--resume_training_at', type=str,
                          help="Location of the pre-trained model to resume training",
                          default=resume_training_at)
      parser.add_argument('--validation_A_dir', type=str,
                          help="validation set for sound source A", default=validation_A_dir_default)
      parser.add_argument('--output_A_dir', type=str,
                          help="output for converted Sound Source A", default=output_A_dir_default)
      parser.add_argument('--validation_B_dir', type=str,
                          help="Validation set for sound source B", default=validation_B_dir_default)
      parser.add_argument('--output_B_dir', type=str,
                          help="Output for converted sound Source B", default=output_B_dir_default)
  
      argv = parser.parse_args()
  
      logf0s_normalization = argv.logf0s_normalization
      mcep_normalization = argv.mcep_normalization
      coded_sps_A_norm = argv.coded_sps_A_norm
      coded_sps_B_norm = argv.coded_sps_B_norm
      logf0s_A_norm = argv.logf0s_A_norm
      logf0s_B_norm = argv.logf0s_B_norm
      model_checkpoint = argv.model_checkpoint
      resume_training_at = argv.resume_training_at
  
      validation_A_dir = argv.validation_A_dir
      output_A_dir = argv.output_A_dir
      validation_B_dir = argv.validation_B_dir
      output_B_dir = argv.output_B_dir
  
      # Check whether following cached files exists
      if not os.path.exists(logf0s_normalization) or not os.path.exists(mcep_normalization):
          print(
              "Cached files do not exist, please run the program preprocess_training.py first")
  
      cycleGAN = CycleGANTraining(logf0s_normalization=logf0s_normalization,
                                  mcep_normalization=mcep_normalization,
                                  coded_sps_A_norm=coded_sps_A_norm,
                                  coded_sps_B_norm=coded_sps_B_norm,
                                  logf0s_A_norm=logf0s_A_norm,
                                  logf0s_B_norm=logf0s_B_norm,
                                  model_checkpoint=model_checkpoint,
                                  validation_A_dir=validation_A_dir,
                                  output_A_dir=output_A_dir,
                                  validation_B_dir=validation_B_dir,
                                  output_B_dir=output_B_dir,
                                  restart_training_at=resume_training_at)
      cycleGAN.train()
  