v1

monotonic_align/core.pyx

+cimport cython
+from cython.parallel import prange
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cdef void maximum_path_each(int[:,::1] path, float[:,::1] value, int t_y, int t_x, float max_neg_val=-1e9) nogil:
+  cdef int x
+  cdef int y
+  cdef float v_prev
+  cdef float v_cur
+  cdef float tmp
+  cdef int index = t_x - 1
+
+  for y in range(t_y):
+    for x in range(max(0, t_x + y - t_y), min(t_x, y + 1)):
+      if x == y:
+        v_cur = max_neg_val
+      else:
+        v_cur = value[y-1, x]
+      if x == 0:
+        if y == 0:
+          v_prev = 0.
+        else:
+          v_prev = max_neg_val
+      else:
+        v_prev = value[y-1, x-1]
+      value[y, x] += max(v_prev, v_cur)
+
+  for y in range(t_y - 1, -1, -1):
+    path[y, index] = 1
+    if index != 0 and (index == y or value[y-1, index] < value[y-1, index-1]):
+      index = index - 1
+
+
+@cython.boundscheck(False)
+@cython.wraparound(False)
+cpdef void maximum_path_c(int[:,:,::1] paths, float[:,:,::1] values, int[::1] t_ys, int[::1] t_xs) nogil:
+  cdef int b = paths.shape[0]
+  cdef int i
+  for i in prange(b, nogil=True):
+    maximum_path_each(paths[i], values[i], t_ys[i], t_xs[i])

monotonic_align/setup.py

+from distutils.core import setup
+from Cython.Build import cythonize
+import numpy
+
+setup(
+  name = 'monotonic_align',
+  ext_modules = cythonize("core.pyx"),
+  include_dirs=[numpy.get_include()]
+)

preprocess.py

+import argparse
+import text
+from utils import load_filepaths_and_text
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--out_extension", default="cleaned")
+  parser.add_argument("--text_index", default=1, type=int)
+  parser.add_argument("--filelists", nargs="+", default=["filelists/ljs_audio_text_val_filelist.txt", "filelists/ljs_audio_text_test_filelist.txt"])
+  parser.add_argument("--text_cleaners", nargs="+", default=["english_cleaners2"])
+
+  args = parser.parse_args()
+    
+
+  for filelist in args.filelists:
+    print("START:", filelist)
+    filepaths_and_text = load_filepaths_and_text(filelist)
+    for i in range(len(filepaths_and_text)):
+      original_text = filepaths_and_text[i][args.text_index]
+      cleaned_text = text._clean_text(original_text, args.text_cleaners)
+      filepaths_and_text[i][args.text_index] = cleaned_text
+
+    new_filelist = filelist + "." + args.out_extension
+    with open(new_filelist, "w", encoding="utf-8") as f:
+      f.writelines(["|".join(x) + "\n" for x in filepaths_and_text])

requirements.txt

+Cython
+librosa==0.9.1
+matplotlib
+numpy==1.24.0
+phonemizer
+scipy
+tensorboard
+Unidecode

run_multi.sh

+source /opt/dtk-24.04.1/env.sh
+export HIP_VISIBLE_DEVICES=6,7
+export HSA_FORCE_FINE_GRAIN_PCIE=1
+
+#export HIP_LOG_LEVEL=1 
+#export HIP_MODULE_MASK=0x7FFFFFFF 
+#export HIP_ENABLE_LIST="hip"
+#export MIOPEN_ENABLE_LOGGING_CMD=1
+#export MIOPEN_LOG_LEVEL=5
+#export MIOPEN_DEBUG_HIPNN_CONV_S868=0
+export MIOPEN_DEBUG_HIPNN_CONV_S390=0
+#export MIOPEN_PRECISION_FP32_FP32_FP32_TF32_FP32=1
+python train.py -c configs/ljs_base.json -m ljs_base

run_single.sh

+source /opt/dtk-24.04.1/env.sh
+export HIP_VISIBLE_DEVICES=6
+
+export MIOPEN_DEBUG_HIPNN_CONV_S390=0
+python train.py -c configs/ljs_base.json -m ljs_base

text/LICENSE

+Copyright (c) 2017 Keith Ito
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

text/__init__.py

+""" from https://github.com/keithito/tacotron """
+from text import cleaners
+from text.symbols import symbols
+
+
+# Mappings from symbol to numeric ID and vice versa:
+_symbol_to_id = {s: i for i, s in enumerate(symbols)}
+_id_to_symbol = {i: s for i, s in enumerate(symbols)}
+
+
+def text_to_sequence(text, cleaner_names):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+      cleaner_names: names of the cleaner functions to run the text through
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  sequence = []
+
+  clean_text = _clean_text(text, cleaner_names)
+  for symbol in clean_text:
+    symbol_id = _symbol_to_id[symbol]
+    sequence += [symbol_id]
+  return sequence
+
+
+def cleaned_text_to_sequence(cleaned_text):
+  '''Converts a string of text to a sequence of IDs corresponding to the symbols in the text.
+    Args:
+      text: string to convert to a sequence
+    Returns:
+      List of integers corresponding to the symbols in the text
+  '''
+  sequence = [_symbol_to_id[symbol] for symbol in cleaned_text]
+  return sequence
+
+
+def sequence_to_text(sequence):
+  '''Converts a sequence of IDs back to a string'''
+  result = ''
+  for symbol_id in sequence:
+    s = _id_to_symbol[symbol_id]
+    result += s
+  return result
+
+
+def _clean_text(text, cleaner_names):
+  for name in cleaner_names:
+    cleaner = getattr(cleaners, name)
+    if not cleaner:
+      raise Exception('Unknown cleaner: %s' % name)
+    text = cleaner(text)
+  return text

text/cleaners.py

+""" from https://github.com/keithito/tacotron """
+
+'''
+Cleaners are transformations that run over the input text at both training and eval time.
+
+Cleaners can be selected by passing a comma-delimited list of cleaner names as the "cleaners"
+hyperparameter. Some cleaners are English-specific. You'll typically want to use:
+  1. "english_cleaners" for English text
+  2. "transliteration_cleaners" for non-English text that can be transliterated to ASCII using
+     the Unidecode library (https://pypi.python.org/pypi/Unidecode)
+  3. "basic_cleaners" if you do not want to transliterate (in this case, you should also update
+     the symbols in symbols.py to match your data).
+'''
+
+import re
+from unidecode import unidecode
+from phonemizer import phonemize
+
+
+# Regular expression matching whitespace:
+_whitespace_re = re.compile(r'\s+')
+
+# List of (regular expression, replacement) pairs for abbreviations:
+_abbreviations = [(re.compile('\\b%s\\.' % x[0], re.IGNORECASE), x[1]) for x in [
+  ('mrs', 'misess'),
+  ('mr', 'mister'),
+  ('dr', 'doctor'),
+  ('st', 'saint'),
+  ('co', 'company'),
+  ('jr', 'junior'),
+  ('maj', 'major'),
+  ('gen', 'general'),
+  ('drs', 'doctors'),
+  ('rev', 'reverend'),
+  ('lt', 'lieutenant'),
+  ('hon', 'honorable'),
+  ('sgt', 'sergeant'),
+  ('capt', 'captain'),
+  ('esq', 'esquire'),
+  ('ltd', 'limited'),
+  ('col', 'colonel'),
+  ('ft', 'fort'),
+]]
+
+
+def expand_abbreviations(text):
+  for regex, replacement in _abbreviations:
+    text = re.sub(regex, replacement, text)
+  return text
+
+
+def expand_numbers(text):
+  return normalize_numbers(text)
+
+
+def lowercase(text):
+  return text.lower()
+
+
+def collapse_whitespace(text):
+  return re.sub(_whitespace_re, ' ', text)
+
+
+def convert_to_ascii(text):
+  return unidecode(text)
+
+
+def basic_cleaners(text):
+  '''Basic pipeline that lowercases and collapses whitespace without transliteration.'''
+  text = lowercase(text)
+  text = collapse_whitespace(text)
+  return text
+
+
+def transliteration_cleaners(text):
+  '''Pipeline for non-English text that transliterates to ASCII.'''
+  text = convert_to_ascii(text)
+  text = lowercase(text)
+  text = collapse_whitespace(text)
+  return text
+
+
+def english_cleaners(text):
+  '''Pipeline for English text, including abbreviation expansion.'''
+  text = convert_to_ascii(text)
+  text = lowercase(text)
+  text = expand_abbreviations(text)
+  phonemes = phonemize(text, language='en-us', backend='espeak', strip=True)
+  phonemes = collapse_whitespace(phonemes)
+  return phonemes
+
+
+def english_cleaners2(text):
+  '''Pipeline for English text, including abbreviation expansion. + punctuation + stress'''
+  text = convert_to_ascii(text)
+  text = lowercase(text)
+  text = expand_abbreviations(text)
+  phonemes = phonemize(text, language='en-us', backend='espeak', strip=True, preserve_punctuation=True, with_stress=True)
+  phonemes = collapse_whitespace(phonemes)
+  return phonemes

text/symbols.py

+""" from https://github.com/keithito/tacotron """
+
+'''
+Defines the set of symbols used in text input to the model.
+'''
+_pad        = '_'
+_punctuation = ';:,.!?¡¿—…"«»“” '
+_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
+_letters_ipa = "ɑɐɒæɓʙβɔɕçɗɖðʤəɘɚɛɜɝɞɟʄɡɠɢʛɦɧħɥʜɨɪʝɭɬɫɮʟɱɯɰŋɳɲɴøɵɸθœɶʘɹɺɾɻʀʁɽʂʃʈʧʉʊʋⱱʌɣɤʍχʎʏʑʐʒʔʡʕʢǀǁǂǃˈˌːˑʼʴʰʱʲʷˠˤ˞↓↑→↗↘'̩'ᵻ"
+
+
+# Export all symbols:
+symbols = [_pad] + list(_punctuation) + list(_letters) + list(_letters_ipa)
+
+# Special symbol ids
+SPACE_ID = symbols.index(" ")

train.py

+import os
+import json
+import argparse
+import itertools
+import math
+import torch
+from torch import nn, optim
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+import torch.multiprocessing as mp
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.cuda.amp import autocast, GradScaler
+
+import commons
+import utils
+from data_utils import (
+  TextAudioLoader,
+  TextAudioCollate,
+  DistributedBucketSampler
+)
+from models import (
+  SynthesizerTrn,
+  MultiPeriodDiscriminator,
+)
+from losses import (
+  generator_loss,
+  discriminator_loss,
+  feature_loss,
+  kl_loss
+)
+from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
+from text.symbols import symbols
+import time
+
+global_step = 0
+
+
+def main():
+  """Assume Single Node Multi GPUs Training Only"""
+  assert torch.cuda.is_available(), "CPU training is not allowed."
+
+  n_gpus = torch.cuda.device_count()
+  os.environ['MASTER_ADDR'] = 'localhost'
+  os.environ['MASTER_PORT'] = '12354'
+
+  hps = utils.get_hparams()
+  mp.spawn(run, nprocs=n_gpus, args=(n_gpus, hps,))
+
+
+def run(rank, n_gpus, hps):
+  global global_step
+  if rank == 0:
+    logger = utils.get_logger(hps.model_dir)
+    logger.info(hps)
+    utils.check_git_hash(hps.model_dir)
+    writer = SummaryWriter(log_dir=hps.model_dir)
+    writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
+
+  dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)
+  torch.manual_seed(hps.train.seed)
+  torch.cuda.set_device(rank)
+
+  train_dataset = TextAudioLoader(hps.data.training_files, hps.data)
+  train_sampler = DistributedBucketSampler(
+      train_dataset,
+      hps.train.batch_size,
+      [32,300,400,500,600,700,800,900,1000],
+      num_replicas=n_gpus,
+      rank=rank,
+      shuffle=True)
+  collate_fn = TextAudioCollate()
+  train_loader = DataLoader(train_dataset, num_workers=8, shuffle=False, pin_memory=True,
+      collate_fn=collate_fn, batch_sampler=train_sampler)
+  if rank == 0:
+    eval_dataset = TextAudioLoader(hps.data.validation_files, hps.data)
+    eval_loader = DataLoader(eval_dataset, num_workers=8, shuffle=False,
+        batch_size=hps.train.batch_size, pin_memory=True,
+        drop_last=False, collate_fn=collate_fn)
+
+  net_g = SynthesizerTrn(
+      len(symbols),
+      hps.data.filter_length // 2 + 1,
+      hps.train.segment_size // hps.data.hop_length,
+      **hps.model).cuda(rank)
+  net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(rank)
+  if rank==0:
+      print(net_g)
+      print(net_d)
+  optim_g = torch.optim.AdamW(
+      net_g.parameters(), 
+      hps.train.learning_rate, 
+      betas=hps.train.betas, 
+      eps=hps.train.eps)
+  optim_d = torch.optim.AdamW(
+      net_d.parameters(),
+      hps.train.learning_rate, 
+      betas=hps.train.betas, 
+      eps=hps.train.eps)
+  net_g = DDP(net_g, device_ids=[rank])
+  net_d = DDP(net_d, device_ids=[rank])
+
+  try:
+    _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g)
+    _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, optim_d)
+    global_step = (epoch_str - 1) * len(train_loader)
+  except:
+    epoch_str = 1
+    global_step = 0
+
+  scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str-2)
+  scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str-2)
+
+  scaler = GradScaler(enabled=hps.train.fp16_run)
+  total_samples = len(train_loader.dataset)
+  for epoch in range(epoch_str, hps.train.epochs + 1):
+    if rank==0:
+      logger.info('====> Epoch: {} | Total Samples: {}'.format(epoch, total_samples))
+      train_and_evaluate(rank, epoch, hps, [net_g, net_d], [optim_g, optim_d], [scheduler_g, scheduler_d], scaler, [train_loader, eval_loader], logger, [writer, writer_eval])
+    else:
+      train_and_evaluate(rank, epoch, hps, [net_g, net_d], [optim_g, optim_d], [scheduler_g, scheduler_d], scaler, [train_loader, None], None, None)
+    scheduler_g.step()
+    scheduler_d.step()
+
+
+def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loaders, logger, writers):
+  net_g, net_d = nets
+  optim_g, optim_d = optims
+  scheduler_g, scheduler_d = schedulers
+  train_loader, eval_loader = loaders
+  if writers is not None:
+    writer, writer_eval = writers
+
+  train_loader.batch_sampler.set_epoch(epoch)
+  global global_step
+
+  net_g.train()
+  net_d.train()
+  start_time = time.time()
+  for batch_idx, (x, x_lengths, spec, spec_lengths, y, y_lengths) in enumerate(train_loader):
+    x, x_lengths = x.cuda(rank, non_blocking=True), x_lengths.cuda(rank, non_blocking=True)
+    spec, spec_lengths = spec.cuda(rank, non_blocking=True), spec_lengths.cuda(rank, non_blocking=True)
+    y, y_lengths = y.cuda(rank, non_blocking=True), y_lengths.cuda(rank, non_blocking=True)
+
+    with autocast(enabled=hps.train.fp16_run):
+      y_hat, l_length, attn, ids_slice, x_mask, z_mask,\
+      (z, z_p, m_p, logs_p, m_q, logs_q) = net_g(x, x_lengths, spec, spec_lengths)
+
+      mel = spec_to_mel_torch(
+          spec, 
+          hps.data.filter_length, 
+          hps.data.n_mel_channels, 
+          hps.data.sampling_rate,
+          hps.data.mel_fmin, 
+          hps.data.mel_fmax)
+      y_mel = commons.slice_segments(mel, ids_slice, hps.train.segment_size // hps.data.hop_length)
+      y_hat_mel = mel_spectrogram_torch(
+          y_hat.squeeze(1), 
+          hps.data.filter_length, 
+          hps.data.n_mel_channels, 
+          hps.data.sampling_rate, 
+          hps.data.hop_length, 
+          hps.data.win_length, 
+          hps.data.mel_fmin, 
+          hps.data.mel_fmax
+      )
+
+      y = commons.slice_segments(y, ids_slice * hps.data.hop_length, hps.train.segment_size) # slice 
+
+      # Discriminator
+      y_d_hat_r, y_d_hat_g, _, _ = net_d(y, y_hat.detach())
+      with autocast(enabled=False):
+        loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
+        loss_disc_all = loss_disc
+    optim_d.zero_grad()
+    scaler.scale(loss_disc_all).backward()
+    scaler.unscale_(optim_d)
+    grad_norm_d = commons.clip_grad_value_(net_d.parameters(), None)
+    scaler.step(optim_d)
+
+    with autocast(enabled=hps.train.fp16_run):
+      # Generator
+      y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y, y_hat)
+      with autocast(enabled=False):
+        loss_dur = torch.sum(l_length.float())
+        loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel
+        loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.c_kl
+
+        loss_fm = feature_loss(fmap_r, fmap_g)
+        loss_gen, losses_gen = generator_loss(y_d_hat_g)
+        loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl
+    optim_g.zero_grad()
+    scaler.scale(loss_gen_all).backward()
+    scaler.unscale_(optim_g)
+    grad_norm_g = commons.clip_grad_value_(net_g.parameters(), None)
+    scaler.step(optim_g)
+    scaler.update()
+
+    if rank==0:
+      if global_step % hps.train.log_interval == 0:
+        lr = optim_g.param_groups[0]['lr']
+        losses = [loss_disc, loss_gen, loss_fm, loss_mel, loss_dur, loss_kl]
+        logger.info('Train Epoch: {} [{:.0f}%]'.format(
+          epoch,
+          100. * batch_idx / len(train_loader)))
+        logger.info([x.item() for x in losses] + [global_step, lr])
+        
+        scalar_dict = {"loss/g/total": loss_gen_all, "loss/d/total": loss_disc_all, "learning_rate": lr, "grad_norm_d": grad_norm_d, "grad_norm_g": grad_norm_g}
+        scalar_dict.update({"loss/g/fm": loss_fm, "loss/g/mel": loss_mel, "loss/g/dur": loss_dur, "loss/g/kl": loss_kl})
+
+        scalar_dict.update({"loss/g/{}".format(i): v for i, v in enumerate(losses_gen)})
+        scalar_dict.update({"loss/d_r/{}".format(i): v for i, v in enumerate(losses_disc_r)})
+        scalar_dict.update({"loss/d_g/{}".format(i): v for i, v in enumerate(losses_disc_g)})
+        image_dict = { 
+            "slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
+            "slice/mel_gen": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()), 
+            "all/mel": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()),
+            "all/attn": utils.plot_alignment_to_numpy(attn[0,0].data.cpu().numpy())
+        }
+        utils.summarize(
+          writer=writer,
+          global_step=global_step, 
+          images=image_dict,
+          scalars=scalar_dict)
+
+      if global_step % hps.train.eval_interval == 0:
+        evaluate(hps, net_g, eval_loader, writer_eval)
+        utils.save_checkpoint(net_g, optim_g, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "G_{}.pth".format(global_step)))
+        utils.save_checkpoint(net_d, optim_d, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "D_{}.pth".format(global_step)))
+    global_step += 1
+  end_time = time.time()
+  epoch_time = end_time - start_time
+  fps = len(train_loader.dataset) / epoch_time
+  if rank == 0:
+    logger.info('====> Epoch: {}'.format(epoch))
+    logger.info('====> FPS: {}'.format(fps))
+ 
+def evaluate(hps, generator, eval_loader, writer_eval):
+    generator.eval()
+    with torch.no_grad():
+      for batch_idx, (x, x_lengths, spec, spec_lengths, y, y_lengths) in enumerate(eval_loader):
+        x, x_lengths = x.cuda(0), x_lengths.cuda(0)
+        spec, spec_lengths = spec.cuda(0), spec_lengths.cuda(0)
+        y, y_lengths = y.cuda(0), y_lengths.cuda(0)
+
+        # remove else
+        x = x[:1]
+        x_lengths = x_lengths[:1]
+        spec = spec[:1]
+        spec_lengths = spec_lengths[:1]
+        y = y[:1]
+        y_lengths = y_lengths[:1]
+        break
+      y_hat, attn, mask, *_ = generator.module.infer(x, x_lengths, max_len=1000)
+      y_hat_lengths = mask.sum([1,2]).long() * hps.data.hop_length
+
+      mel = spec_to_mel_torch(
+        spec, 
+        hps.data.filter_length, 
+        hps.data.n_mel_channels, 
+        hps.data.sampling_rate,
+        hps.data.mel_fmin, 
+        hps.data.mel_fmax)
+      y_hat_mel = mel_spectrogram_torch(
+        y_hat.squeeze(1).float(),
+        hps.data.filter_length,
+        hps.data.n_mel_channels,
+        hps.data.sampling_rate,
+        hps.data.hop_length,
+        hps.data.win_length,
+        hps.data.mel_fmin,
+        hps.data.mel_fmax
+      )
+    image_dict = {
+      "gen/mel": utils.plot_spectrogram_to_numpy(y_hat_mel[0].cpu().numpy())
+    }
+    audio_dict = {
+      "gen/audio": y_hat[0,:,:y_hat_lengths[0]]
+    }
+    if global_step == 0:
+      image_dict.update({"gt/mel": utils.plot_spectrogram_to_numpy(mel[0].cpu().numpy())})
+      audio_dict.update({"gt/audio": y[0,:,:y_lengths[0]]})
+
+    utils.summarize(
+      writer=writer_eval,
+      global_step=global_step, 
+      images=image_dict,
+      audios=audio_dict,
+      audio_sampling_rate=hps.data.sampling_rate
+    )
+    generator.train()
+
+                           
+if __name__ == "__main__":
+  main()

train_ms.py

+import os
+import json
+import argparse
+import itertools
+import math
+import torch
+from torch import nn, optim
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+import torch.multiprocessing as mp
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.cuda.amp import autocast, GradScaler
+
+import commons
+import utils
+from data_utils import (
+  TextAudioSpeakerLoader,
+  TextAudioSpeakerCollate,
+  DistributedBucketSampler
+)
+from models import (
+  SynthesizerTrn,
+  MultiPeriodDiscriminator,
+)
+from losses import (
+  generator_loss,
+  discriminator_loss,
+  feature_loss,
+  kl_loss
+)
+from mel_processing import mel_spectrogram_torch, spec_to_mel_torch
+from text.symbols import symbols
+
+
+torch.backends.cudnn.benchmark = True
+global_step = 0
+
+
+def main():
+  """Assume Single Node Multi GPUs Training Only"""
+  assert torch.cuda.is_available(), "CPU training is not allowed."
+
+  n_gpus = torch.cuda.device_count()
+  os.environ['MASTER_ADDR'] = 'localhost'
+  os.environ['MASTER_PORT'] = '80000'
+
+  hps = utils.get_hparams()
+  mp.spawn(run, nprocs=n_gpus, args=(n_gpus, hps,))
+
+
+def run(rank, n_gpus, hps):
+  global global_step
+  if rank == 0:
+    logger = utils.get_logger(hps.model_dir)
+    logger.info(hps)
+    utils.check_git_hash(hps.model_dir)
+    writer = SummaryWriter(log_dir=hps.model_dir)
+    writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
+
+  dist.init_process_group(backend='nccl', init_method='env://', world_size=n_gpus, rank=rank)
+  torch.manual_seed(hps.train.seed)
+  torch.cuda.set_device(rank)
+
+  train_dataset = TextAudioSpeakerLoader(hps.data.training_files, hps.data)
+  train_sampler = DistributedBucketSampler(
+      train_dataset,
+      hps.train.batch_size,
+      [32,300,400,500,600,700,800,900,1000],
+      num_replicas=n_gpus,
+      rank=rank,
+      shuffle=True)
+  collate_fn = TextAudioSpeakerCollate()
+  train_loader = DataLoader(train_dataset, num_workers=8, shuffle=False, pin_memory=True,
+      collate_fn=collate_fn, batch_sampler=train_sampler)
+  if rank == 0:
+    eval_dataset = TextAudioSpeakerLoader(hps.data.validation_files, hps.data)
+    eval_loader = DataLoader(eval_dataset, num_workers=8, shuffle=False,
+        batch_size=hps.train.batch_size, pin_memory=True,
+        drop_last=False, collate_fn=collate_fn)
+
+  net_g = SynthesizerTrn(
+      len(symbols),
+      hps.data.filter_length // 2 + 1,
+      hps.train.segment_size // hps.data.hop_length,
+      n_speakers=hps.data.n_speakers,
+      **hps.model).cuda(rank)
+  net_d = MultiPeriodDiscriminator(hps.model.use_spectral_norm).cuda(rank)
+  optim_g = torch.optim.AdamW(
+      net_g.parameters(), 
+      hps.train.learning_rate, 
+      betas=hps.train.betas, 
+      eps=hps.train.eps)
+  optim_d = torch.optim.AdamW(
+      net_d.parameters(),
+      hps.train.learning_rate, 
+      betas=hps.train.betas, 
+      eps=hps.train.eps)
+  net_g = DDP(net_g, device_ids=[rank])
+  net_d = DDP(net_d, device_ids=[rank])
+
+  try:
+    _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "G_*.pth"), net_g, optim_g)
+    _, _, _, epoch_str = utils.load_checkpoint(utils.latest_checkpoint_path(hps.model_dir, "D_*.pth"), net_d, optim_d)
+    global_step = (epoch_str - 1) * len(train_loader)
+  except:
+    epoch_str = 1
+    global_step = 0
+
+  scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=hps.train.lr_decay, last_epoch=epoch_str-2)
+  scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=hps.train.lr_decay, last_epoch=epoch_str-2)
+
+  scaler = GradScaler(enabled=hps.train.fp16_run)
+
+  for epoch in range(epoch_str, hps.train.epochs + 1):
+    if rank==0:
+      train_and_evaluate(rank, epoch, hps, [net_g, net_d], [optim_g, optim_d], [scheduler_g, scheduler_d], scaler, [train_loader, eval_loader], logger, [writer, writer_eval])
+    else:
+      train_and_evaluate(rank, epoch, hps, [net_g, net_d], [optim_g, optim_d], [scheduler_g, scheduler_d], scaler, [train_loader, None], None, None)
+    scheduler_g.step()
+    scheduler_d.step()
+
+
+def train_and_evaluate(rank, epoch, hps, nets, optims, schedulers, scaler, loaders, logger, writers):
+  net_g, net_d = nets
+  optim_g, optim_d = optims
+  scheduler_g, scheduler_d = schedulers
+  train_loader, eval_loader = loaders
+  if writers is not None:
+    writer, writer_eval = writers
+
+  train_loader.batch_sampler.set_epoch(epoch)
+  global global_step
+
+  net_g.train()
+  net_d.train()
+  for batch_idx, (x, x_lengths, spec, spec_lengths, y, y_lengths, speakers) in enumerate(train_loader):
+    x, x_lengths = x.cuda(rank, non_blocking=True), x_lengths.cuda(rank, non_blocking=True)
+    spec, spec_lengths = spec.cuda(rank, non_blocking=True), spec_lengths.cuda(rank, non_blocking=True)
+    y, y_lengths = y.cuda(rank, non_blocking=True), y_lengths.cuda(rank, non_blocking=True)
+    speakers = speakers.cuda(rank, non_blocking=True)
+
+    with autocast(enabled=hps.train.fp16_run):
+      y_hat, l_length, attn, ids_slice, x_mask, z_mask,\
+      (z, z_p, m_p, logs_p, m_q, logs_q) = net_g(x, x_lengths, spec, spec_lengths, speakers)
+
+      mel = spec_to_mel_torch(
+          spec, 
+          hps.data.filter_length, 
+          hps.data.n_mel_channels, 
+          hps.data.sampling_rate,
+          hps.data.mel_fmin, 
+          hps.data.mel_fmax)
+      y_mel = commons.slice_segments(mel, ids_slice, hps.train.segment_size // hps.data.hop_length)
+      y_hat_mel = mel_spectrogram_torch(
+          y_hat.squeeze(1), 
+          hps.data.filter_length, 
+          hps.data.n_mel_channels, 
+          hps.data.sampling_rate, 
+          hps.data.hop_length, 
+          hps.data.win_length, 
+          hps.data.mel_fmin, 
+          hps.data.mel_fmax
+      )
+
+      y = commons.slice_segments(y, ids_slice * hps.data.hop_length, hps.train.segment_size) # slice 
+
+      # Discriminator
+      y_d_hat_r, y_d_hat_g, _, _ = net_d(y, y_hat.detach())
+      with autocast(enabled=False):
+        loss_disc, losses_disc_r, losses_disc_g = discriminator_loss(y_d_hat_r, y_d_hat_g)
+        loss_disc_all = loss_disc
+    optim_d.zero_grad()
+    scaler.scale(loss_disc_all).backward()
+    scaler.unscale_(optim_d)
+    grad_norm_d = commons.clip_grad_value_(net_d.parameters(), None)
+    scaler.step(optim_d)
+
+    with autocast(enabled=hps.train.fp16_run):
+      # Generator
+      y_d_hat_r, y_d_hat_g, fmap_r, fmap_g = net_d(y, y_hat)
+      with autocast(enabled=False):
+        loss_dur = torch.sum(l_length.float())
+        loss_mel = F.l1_loss(y_mel, y_hat_mel) * hps.train.c_mel
+        loss_kl = kl_loss(z_p, logs_q, m_p, logs_p, z_mask) * hps.train.c_kl
+
+        loss_fm = feature_loss(fmap_r, fmap_g)
+        loss_gen, losses_gen = generator_loss(y_d_hat_g)
+        loss_gen_all = loss_gen + loss_fm + loss_mel + loss_dur + loss_kl
+    optim_g.zero_grad()
+    scaler.scale(loss_gen_all).backward()
+    scaler.unscale_(optim_g)
+    grad_norm_g = commons.clip_grad_value_(net_g.parameters(), None)
+    scaler.step(optim_g)
+    scaler.update()
+
+    if rank==0:
+      if global_step % hps.train.log_interval == 0:
+        lr = optim_g.param_groups[0]['lr']
+        losses = [loss_disc, loss_gen, loss_fm, loss_mel, loss_dur, loss_kl]
+        logger.info('Train Epoch: {} [{:.0f}%]'.format(
+          epoch,
+          100. * batch_idx / len(train_loader)))
+        logger.info([x.item() for x in losses] + [global_step, lr])
+        
+        scalar_dict = {"loss/g/total": loss_gen_all, "loss/d/total": loss_disc_all, "learning_rate": lr, "grad_norm_d": grad_norm_d, "grad_norm_g": grad_norm_g}
+        scalar_dict.update({"loss/g/fm": loss_fm, "loss/g/mel": loss_mel, "loss/g/dur": loss_dur, "loss/g/kl": loss_kl})
+
+        scalar_dict.update({"loss/g/{}".format(i): v for i, v in enumerate(losses_gen)})
+        scalar_dict.update({"loss/d_r/{}".format(i): v for i, v in enumerate(losses_disc_r)})
+        scalar_dict.update({"loss/d_g/{}".format(i): v for i, v in enumerate(losses_disc_g)})
+        image_dict = { 
+            "slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
+            "slice/mel_gen": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()), 
+            "all/mel": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()),
+            "all/attn": utils.plot_alignment_to_numpy(attn[0,0].data.cpu().numpy())
+        }
+        utils.summarize(
+          writer=writer,
+          global_step=global_step, 
+          images=image_dict,
+          scalars=scalar_dict)
+
+      if global_step % hps.train.eval_interval == 0:
+        evaluate(hps, net_g, eval_loader, writer_eval)
+        utils.save_checkpoint(net_g, optim_g, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "G_{}.pth".format(global_step)))
+        utils.save_checkpoint(net_d, optim_d, hps.train.learning_rate, epoch, os.path.join(hps.model_dir, "D_{}.pth".format(global_step)))
+    global_step += 1
+  
+  if rank == 0:
+    logger.info('====> Epoch: {}'.format(epoch))
+
+ 
+def evaluate(hps, generator, eval_loader, writer_eval):
+    generator.eval()
+    with torch.no_grad():
+      for batch_idx, (x, x_lengths, spec, spec_lengths, y, y_lengths, speakers) in enumerate(eval_loader):
+        x, x_lengths = x.cuda(0), x_lengths.cuda(0)
+        spec, spec_lengths = spec.cuda(0), spec_lengths.cuda(0)
+        y, y_lengths = y.cuda(0), y_lengths.cuda(0)
+        speakers = speakers.cuda(0)
+
+        # remove else
+        x = x[:1]
+        x_lengths = x_lengths[:1]
+        spec = spec[:1]
+        spec_lengths = spec_lengths[:1]
+        y = y[:1]
+        y_lengths = y_lengths[:1]
+        speakers = speakers[:1]
+        break
+      y_hat, attn, mask, *_ = generator.module.infer(x, x_lengths, speakers, max_len=1000)
+      y_hat_lengths = mask.sum([1,2]).long() * hps.data.hop_length
+
+      mel = spec_to_mel_torch(
+        spec, 
+        hps.data.filter_length, 
+        hps.data.n_mel_channels, 
+        hps.data.sampling_rate,
+        hps.data.mel_fmin, 
+        hps.data.mel_fmax)
+      y_hat_mel = mel_spectrogram_torch(
+        y_hat.squeeze(1).float(),
+        hps.data.filter_length,
+        hps.data.n_mel_channels,
+        hps.data.sampling_rate,
+        hps.data.hop_length,
+        hps.data.win_length,
+        hps.data.mel_fmin,
+        hps.data.mel_fmax
+      )
+    image_dict = {
+      "gen/mel": utils.plot_spectrogram_to_numpy(y_hat_mel[0].cpu().numpy())
+    }
+    audio_dict = {
+      "gen/audio": y_hat[0,:,:y_hat_lengths[0]]
+    }
+    if global_step == 0:
+      image_dict.update({"gt/mel": utils.plot_spectrogram_to_numpy(mel[0].cpu().numpy())})
+      audio_dict.update({"gt/audio": y[0,:,:y_lengths[0]]})
+
+    utils.summarize(
+      writer=writer_eval,
+      global_step=global_step, 
+      images=image_dict,
+      audios=audio_dict,
+      audio_sampling_rate=hps.data.sampling_rate
+    )
+    generator.train()
+
+                           
+if __name__ == "__main__":
+  main()

transforms.py

+import torch
+from torch.nn import functional as F
+
+import numpy as np
+
+
+DEFAULT_MIN_BIN_WIDTH = 1e-3
+DEFAULT_MIN_BIN_HEIGHT = 1e-3
+DEFAULT_MIN_DERIVATIVE = 1e-3
+
+
+def piecewise_rational_quadratic_transform(inputs, 
+                                           unnormalized_widths,
+                                           unnormalized_heights,
+                                           unnormalized_derivatives,
+                                           inverse=False,
+                                           tails=None, 
+                                           tail_bound=1.,
+                                           min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+                                           min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+                                           min_derivative=DEFAULT_MIN_DERIVATIVE):
+
+    if tails is None:
+        spline_fn = rational_quadratic_spline
+        spline_kwargs = {}
+    else:
+        spline_fn = unconstrained_rational_quadratic_spline
+        spline_kwargs = {
+            'tails': tails,
+            'tail_bound': tail_bound
+        }
+
+    outputs, logabsdet = spline_fn(
+            inputs=inputs,
+            unnormalized_widths=unnormalized_widths,
+            unnormalized_heights=unnormalized_heights,
+            unnormalized_derivatives=unnormalized_derivatives,
+            inverse=inverse,
+            min_bin_width=min_bin_width,
+            min_bin_height=min_bin_height,
+            min_derivative=min_derivative,
+            **spline_kwargs
+    )
+    return outputs, logabsdet
+
+
+def searchsorted(bin_locations, inputs, eps=1e-6):
+    bin_locations[..., -1] += eps
+    return torch.sum(
+        inputs[..., None] >= bin_locations,
+        dim=-1
+    ) - 1
+
+
+def unconstrained_rational_quadratic_spline(inputs,
+                                            unnormalized_widths,
+                                            unnormalized_heights,
+                                            unnormalized_derivatives,
+                                            inverse=False,
+                                            tails='linear',
+                                            tail_bound=1.,
+                                            min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+                                            min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+                                            min_derivative=DEFAULT_MIN_DERIVATIVE):
+    inside_interval_mask = (inputs >= -tail_bound) & (inputs <= tail_bound)
+    outside_interval_mask = ~inside_interval_mask
+
+    outputs = torch.zeros_like(inputs)
+    logabsdet = torch.zeros_like(inputs)
+
+    if tails == 'linear':
+        unnormalized_derivatives = F.pad(unnormalized_derivatives, pad=(1, 1))
+        constant = np.log(np.exp(1 - min_derivative) - 1)
+        unnormalized_derivatives[..., 0] = constant
+        unnormalized_derivatives[..., -1] = constant
+
+        outputs[outside_interval_mask] = inputs[outside_interval_mask]
+        logabsdet[outside_interval_mask] = 0
+    else:
+        raise RuntimeError('{} tails are not implemented.'.format(tails))
+
+    outputs[inside_interval_mask], logabsdet[inside_interval_mask] = rational_quadratic_spline(
+        inputs=inputs[inside_interval_mask],
+        unnormalized_widths=unnormalized_widths[inside_interval_mask, :],
+        unnormalized_heights=unnormalized_heights[inside_interval_mask, :],
+        unnormalized_derivatives=unnormalized_derivatives[inside_interval_mask, :],
+        inverse=inverse,
+        left=-tail_bound, right=tail_bound, bottom=-tail_bound, top=tail_bound,
+        min_bin_width=min_bin_width,
+        min_bin_height=min_bin_height,
+        min_derivative=min_derivative
+    )
+
+    return outputs, logabsdet
+
+def rational_quadratic_spline(inputs,
+                              unnormalized_widths,
+                              unnormalized_heights,
+                              unnormalized_derivatives,
+                              inverse=False,
+                              left=0., right=1., bottom=0., top=1.,
+                              min_bin_width=DEFAULT_MIN_BIN_WIDTH,
+                              min_bin_height=DEFAULT_MIN_BIN_HEIGHT,
+                              min_derivative=DEFAULT_MIN_DERIVATIVE):
+    if torch.min(inputs) < left or torch.max(inputs) > right:
+        raise ValueError('Input to a transform is not within its domain')
+
+    num_bins = unnormalized_widths.shape[-1]
+
+    if min_bin_width * num_bins > 1.0:
+        raise ValueError('Minimal bin width too large for the number of bins')
+    if min_bin_height * num_bins > 1.0:
+        raise ValueError('Minimal bin height too large for the number of bins')
+
+    widths = F.softmax(unnormalized_widths, dim=-1)
+    widths = min_bin_width + (1 - min_bin_width * num_bins) * widths
+    cumwidths = torch.cumsum(widths, dim=-1)
+    cumwidths = F.pad(cumwidths, pad=(1, 0), mode='constant', value=0.0)
+    cumwidths = (right - left) * cumwidths + left
+    cumwidths[..., 0] = left
+    cumwidths[..., -1] = right
+    widths = cumwidths[..., 1:] - cumwidths[..., :-1]
+
+    derivatives = min_derivative + F.softplus(unnormalized_derivatives)
+
+    heights = F.softmax(unnormalized_heights, dim=-1)
+    heights = min_bin_height + (1 - min_bin_height * num_bins) * heights
+    cumheights = torch.cumsum(heights, dim=-1)
+    cumheights = F.pad(cumheights, pad=(1, 0), mode='constant', value=0.0)
+    cumheights = (top - bottom) * cumheights + bottom
+    cumheights[..., 0] = bottom
+    cumheights[..., -1] = top
+    heights = cumheights[..., 1:] - cumheights[..., :-1]
+
+    if inverse:
+        bin_idx = searchsorted(cumheights, inputs)[..., None]
+    else:
+        bin_idx = searchsorted(cumwidths, inputs)[..., None]
+
+    input_cumwidths = cumwidths.gather(-1, bin_idx)[..., 0]
+    input_bin_widths = widths.gather(-1, bin_idx)[..., 0]
+
+    input_cumheights = cumheights.gather(-1, bin_idx)[..., 0]
+    delta = heights / widths
+    input_delta = delta.gather(-1, bin_idx)[..., 0]
+
+    input_derivatives = derivatives.gather(-1, bin_idx)[..., 0]
+    input_derivatives_plus_one = derivatives[..., 1:].gather(-1, bin_idx)[..., 0]
+
+    input_heights = heights.gather(-1, bin_idx)[..., 0]
+
+    if inverse:
+        a = (((inputs - input_cumheights) * (input_derivatives
+                                             + input_derivatives_plus_one
+                                             - 2 * input_delta)
+              + input_heights * (input_delta - input_derivatives)))
+        b = (input_heights * input_derivatives
+             - (inputs - input_cumheights) * (input_derivatives
+                                              + input_derivatives_plus_one
+                                              - 2 * input_delta))
+        c = - input_delta * (inputs - input_cumheights)
+
+        discriminant = b.pow(2) - 4 * a * c
+        assert (discriminant >= 0).all()
+
+        root = (2 * c) / (-b - torch.sqrt(discriminant))
+        outputs = root * input_bin_widths + input_cumwidths
+
+        theta_one_minus_theta = root * (1 - root)
+        denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+                                     * theta_one_minus_theta)
+        derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * root.pow(2)
+                                                     + 2 * input_delta * theta_one_minus_theta
+                                                     + input_derivatives * (1 - root).pow(2))
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, -logabsdet
+    else:
+        theta = (inputs - input_cumwidths) / input_bin_widths
+        theta_one_minus_theta = theta * (1 - theta)
+
+        numerator = input_heights * (input_delta * theta.pow(2)
+                                     + input_derivatives * theta_one_minus_theta)
+        denominator = input_delta + ((input_derivatives + input_derivatives_plus_one - 2 * input_delta)
+                                     * theta_one_minus_theta)
+        outputs = input_cumheights + numerator / denominator
+
+        derivative_numerator = input_delta.pow(2) * (input_derivatives_plus_one * theta.pow(2)
+                                                     + 2 * input_delta * theta_one_minus_theta
+                                                     + input_derivatives * (1 - theta).pow(2))
+        logabsdet = torch.log(derivative_numerator) - 2 * torch.log(denominator)
+
+        return outputs, logabsdet