assert os.path.isfile(checkpoint_path)

print("Starting model from checkpoint '{}'".format(checkpoint_path))

checkpoint_dict = torch.load(checkpoint_path, map_location='cuda:{}'.format(rank))

if 'model' in checkpoint_dict:

if distributed:

state_dict = {}

for k,v in checkpoint_dict['model'].items():

state_dict['module.{}'.format(k)] = v

model.load_state_dict(state_dict)

else:

model.load_state_dict(checkpoint_dict['model'])

print('Model is loaded!')

if 'discriminator' in checkpoint_dict:

if distributed:

state_dict = {}

for k,v in checkpoint_dict['discriminator'].items():

state_dict['module.{}'.format(k)] = v

discriminator.load_state_dict(state_dict)

else:

discriminator.load_state_dict(checkpoint_dict['discriminator'])

print('Discriminator is loaded!')

if 'G_optim' in checkpoint_dict or 'optimizer' in checkpoint_dict:

if 'optimizer' in checkpoint_dict:

G_optim.load_state_dict(checkpoint_dict['optimizer'])

if 'G_optim' in checkpoint_dict:

G_optim.load_state_dict(checkpoint_dict['G_optim'])

print('G_optim is loaded!')

if 'D_optim' in checkpoint_dict:

D_optim.load_state_dict(checkpoint_dict['D_optim'])

print('D_optim is loaded!')

current_step = checkpoint_dict['step'] + 1

del checkpoint_dict

print('Use GPU: {} for training'.format(rank))

ngpus = args.ngpus

if args.distributed:

torch.cuda.set_device(rank % ngpus)

dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=rank)

# Define model & loss

model = StyleSpeech(c).cuda()

discriminator = Discriminator(c).cuda()

num_param = utils.get_param_num(model)

D_num_param = utils.get_param_num(discriminator)

if rank==0:

print('Number of Meta-StyleSpeech Parameters:', num_param)

print('Number of Discriminator Parameters:', D_num_param)

with open(os.path.join(args.save_path, "model.txt"), "w") as f_log:

f_log.write(str(model))

f_log.write(str(discriminator))

print("Model Has Been Defined")

model_without_ddp = model

discriminator_without_ddp = discriminator

if args.distributed:

c.meta_batch_size = c.meta_batch_size // ngpus

model = nn.parallel.DistributedDataParallel(model, device_ids=[rank])

model_without_ddp = model.module

discriminator = nn.parallel.DistributedDataParallel(discriminator, device_ids=[rank])

discriminator_without_ddp = discriminator.module

# Optimizer

G_optim = torch.optim.Adam(model.parameters(), betas=c.betas, eps=c.eps)

D_optim = torch.optim.Adam(discriminator.parameters(), lr=2e-4, betas=c.betas, eps=c.eps)

# Loss

Loss = model_without_ddp.get_criterion()

adversarial_loss = discriminator_without_ddp.get_criterion()

print("Optimizer and Loss Function Defined.")

# Get dataset

data_loader = prepare_dataloader(args.data_path, "train.txt", batch_size=c.meta_batch_size, meta_learning=True, seed=rank)

print("Data Loader is Prepared")

# Load checkpoint if exists

if args.checkpoint_path is not None:

assert os.path.exists(args.checkpoint_path)

model, discriminator, G_optim, D_optim, current_step = load_checkpoint(

args.checkpoint_path, model, discriminator, G_optim, D_optim, rank, args.distributed)

print("\n---Model Restored at Step {}---\n".format(current_step))

else:

print("\n---Start New Training---\n")

current_step = 0

if rank == 0:

checkpoint_path = os.path.join(args.save_path, 'ckpt')

os.makedirs(checkpoint_path, exist_ok=True)

# scheduled optimizer

G_optim = ScheduledOptim(G_optim, c.decoder_hidden, c.n_warm_up_step, current_step)

# Init logger

if rank == 0:

log_path = os.path.join(args.save_path, 'log')

logger = SummaryWriter(os.path.join(log_path, 'board'))

with open(os.path.join(log_path, "log.txt"), "a") as f_log:

f_log.write("Dataset :{}\n Number of Parameters: {}\n".format(c.dataset, num_param))

# Init synthesis directory

if rank == 0:

synth_path = os.path.join(args.save_path, 'synth')

os.makedirs(synth_path, exist_ok=True)

model.train()

while current_step < args.max_iter:

# Get Training Loader

for idx, batch in enumerate(data_loader):

if current_step == args.max_iter:

break

losses = {}

#### Generator ####

G_optim.zero_grad()

# Get Support Data

sid, text, mel_target, D, log_D, f0, energy, \

src_len, mel_len, max_src_len, max_mel_len = model_without_ddp.parse_batch(batch)

# Support Forward

mel_output, src_output, style_vector, log_duration_output, f0_output, energy_output, src_mask, mel_mask, _ = model(

text, src_len, mel_target, mel_len, D, f0, energy, max_src_len, max_mel_len)

src_target, _, _ = model_without_ddp.variance_adaptor.length_regulator(src_output, D)

# Reconstruction loss

mel_loss, d_loss, f_loss, e_loss = Loss(mel_output, mel_target,

log_duration_output, log_D, f0_output, f0, energy_output, energy, src_len, mel_len)

losses['G_recon'] = mel_loss

losses['d_loss'] = d_loss

losses['f_loss'] = f_loss

losses['e_loss'] = e_loss

#### META LEARNING ####

# Get query text

B = mel_target.shape[0]

perm_idx = torch.randperm(B)

q_text, q_src_len = text[perm_idx], src_len[perm_idx]

# Generate query speech

q_mel_output, q_src_output, q_log_duration_output, \

_, _, q_src_mask, q_mel_mask, q_mel_len = model_without_ddp.inference(style_vector, q_text, q_src_len)

# Legulate length of query src

q_duration_rounded = torch.clamp(torch.round(torch.exp(q_log_duration_output.detach())-1.), min=0)

q_duration = q_duration_rounded.masked_fill(q_src_mask, 0).long()

q_src, _, _ = model_without_ddp.variance_adaptor.length_regulator(q_src_output, q_duration)

# Adverserial loss

t_val, s_val, _= discriminator(q_mel_output, q_src, None, sid, q_mel_mask)

losses['G_GAN_query_t'] = adversarial_loss(t_val, is_real=True)

losses['G_GAN_query_s'] = adversarial_loss(s_val, is_real=True)

# Total generator loss

alpha = 10.0

G_loss = alpha*losses['G_recon'] + losses['d_loss'] + losses['f_loss'] + losses['e_loss'] + \

losses['G_GAN_query_t'] + losses['G_GAN_query_s']

# Backward loss

G_loss.backward()

# Update weights

G_optim.step_and_update_lr()

#### Discriminator ####

D_optim.zero_grad()

# Real

real_t_pred, real_s_pred, cls_loss = discriminator(

mel_target, src_target.detach(), style_vector.detach(), sid, mask=mel_mask)

# Fake

fake_t_pred, fake_s_pred, _ = discriminator(

q_mel_output.detach(), q_src.detach(), None, sid, mask=q_mel_mask)

losses['D_t_loss'] = adversarial_loss(real_t_pred, is_real=True) + adversarial_loss(fake_t_pred, is_real=False)

losses['D_s_loss'] = adversarial_loss(real_s_pred, is_real=True) + adversarial_loss(fake_s_pred, is_real=False)

losses['cls_loss'] = cls_loss

# Total discriminator Loss

D_loss = losses['D_t_loss'] + losses['D_s_loss'] + losses['cls_loss']

# Backward

D_loss.backward()

# Update weights

D_optim.step()

# Print log

if current_step % args.log_step == 0 and current_step != 0 and rank == 0 :

m_l = losses['G_recon'].item()

d_l = losses['d_loss'].item()

f_l = losses['f_loss'].item()

e_l = losses['e_loss'].item()

g_t_l = losses['G_GAN_query_t'].item()

g_s_l = losses['G_GAN_query_s'].item()

d_t_l = losses['D_t_loss'].item() / 2

d_s_l = losses['D_s_loss'].item() / 2

cls_l = losses['cls_loss'].item()

str1 = "Step [{}/{}]:".format(current_step, args.max_iter)

str2 = "Mel Loss: {:.4f},\n" \

"Duration Loss: {:.4f}, F0 Loss: {:.4f}, Energy Loss: {:.4f}\n" \

"T G Loss: {:.4f}, T D Loss: {:.4f}, S G Loss: {:.4f}, S D Loss: {:.4f} \n" \

"cls_Loss: {:.4f};" \

.format(m_l, d_l, f_l, e_l, g_t_l, d_t_l, g_s_l, d_s_l, cls_l)

print(str1 + "\n" + str2 +"\n")

with open(os.path.join(log_path, "log.txt"), "a") as f_log:

f_log.write(str1 + "\n" + str2 +"\n")

logger.add_scalar('Train/mel_loss', m_l, current_step)

logger.add_scalar('Train/duration_loss', d_l, current_step)

logger.add_scalar('Train/f0_loss', f_l, current_step)

logger.add_scalar('Train/energy_loss', e_l, current_step)

logger.add_scalar('Train/G_t_loss', g_t_l, current_step)

logger.add_scalar('Train/D_t_loss', d_t_l, current_step)

logger.add_scalar('Train/G_s_loss', g_s_l, current_step)

logger.add_scalar('Train/D_s_loss', d_s_l, current_step)

logger.add_scalar('Train/cls_loss', cls_l, current_step)

# Save Checkpoint

if current_step % args.save_step == 0 and current_step != 0 and rank == 0:

torch.save({'model': model_without_ddp.state_dict(),

'discriminator': discriminator_without_ddp.state_dict(),

'G_optim': G_optim.state_dict(),'D_optim': D_optim.state_dict(),

'step': current_step},

os.path.join(checkpoint_path, 'checkpoint_{}.pth.tar'.format(current_step)))

print("*** Save Checkpoint ***")

print("Save model at step {}...\n".format(current_step))

if current_step % args.synth_step == 0 and current_step != 0 and rank == 0:

length = mel_len[0].item()

mel_target = mel_target[0, :length].detach().cpu().transpose(0, 1)

mel = mel_output[0, :length].detach().cpu().transpose(0, 1)

q_length = q_mel_len[0].item()

q_mel = q_mel_output[0, :q_length].detach().cpu().transpose(0, 1)

# plotting

utils.plot_data([q_mel.numpy(), mel.numpy(), mel_target.numpy()],

['Query Spectrogram', 'Recon Spectrogram', 'Ground-Truth Spectrogram'], filename=os.path.join(synth_path, 'step_{}.png'.format(current_step)))

print("Synth audios at step {}...\n".format(current_step))

# Evaluate

if current_step % args.eval_step == 0 and current_step != 0 and rank == 0:

model.eval()

with torch.no_grad():

m_l, d_l, f_l, e_l = evaluate(args, model_without_ddp, current_step)

str_v = "*** Validation ***\n" \

"Meta-StyleSpeech Step {},\n" \

"Mel Loss: {}\nDuration Loss:{}\nF0 Loss: {}\nEnergy Loss: {}" \

.format(current_step, m_l, d_l, f_l, e_l)

print(str_v + "\n" )

with open(os.path.join(log_path, "eval.txt"), "a") as f_log:

f_log.write(str_v + "\n")

logger.add_scalar('Validation/mel_loss', m_l, current_step)

logger.add_scalar('Validation/duration_loss', d_l, current_step)

logger.add_scalar('Validation/f0_loss', f_l, current_step)

logger.add_scalar('Validation/energy_loss', e_l, current_step)

model.train()

current_step += 1

if rank == 0:

print("Training Done at Step : {}".format(current_step))

torch.save({'model': model_without_ddp.state_dict(),

'discriminator': discriminator_without_ddp.state_dict(),

'G_optim': G_optim.state_dict(), 'D_optim': D_optim.state_dict(),

'step': current_step},