working integraton with parakeet

.gitignore

@@ -129,4 +129,10 @@ venv.bak/
 dmypy.json
 
 # Pyre type checker
-.pyre/
\ No newline at end of file
+.pyre/
+
+# Shell, vim, and output folder
+*.sh
+*.swp
+runs
+syn_audios

parakeet/data/datacargo.py

@@ -31,6 +31,9 @@ class DataCargo(object):
     
     def __iter__(self):
         return DataIterator(self)
+
+    def __call__(self):
+        return DataIterator(self)
     
     @property
     def _auto_collation(self):

parakeet/models/wavenet/README.md

+# WaveNet-Paddle
\ No newline at end of file

parakeet/models/wavenet/configs/wavenet_ljspeech_single_gaussian.yaml

+valid_size: 16
+train_clip_second: 0.5
+sample_rate: 22050
+fft_window_shift: 256
+fft_window_size: 1024
+fft_size: 2048
+mel_bands: 80
+
+seed: 1
+batch_size: 8
+test_every: 2000
+save_every: 10000
+max_iterations: 2000000
+
+layers: 30
+kernel_width: 2
+dilation_block: [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]
+residual_channels: 128
+skip_channels: 128
+loss_type: mix-gaussian-pdf
+num_mixtures: 1
+log_scale_min: -9.0
+
+conditioner: 
+  filter_sizes: [[32, 3], [32, 3]]
+  upsample_factors: [16, 16]
+
+learning_rate: 0.001
+gradient_max_norm: 100.0
+anneal: 
+  every: 200000
+  rate: 0.5

parakeet/models/wavenet/data.py

+import math
+import os
+import random
+
+import librosa
+import numpy as np
+from paddle import fluid
+
+import utils
+from parakeet.datasets import ljspeech
+from parakeet.data import dataset
+from parakeet.data.sampler import Sampler, BatchSampler, SequentialSampler
+from parakeet.data.datacargo import DataCargo
+
+
+class Dataset(ljspeech.LJSpeech):
+    def __init__(self, config):
+        super(Dataset, self).__init__(config.root)
+        self.config = config
+        self.fft_window_shift = config.fft_window_shift
+        # Calculate context frames.
+        frames_per_second = config.sample_rate // self.fft_window_shift
+        train_clip_frames = int(math.ceil(
+            config.train_clip_second * frames_per_second))
+        context_frames = config.context_size // self.fft_window_shift
+        self.num_frames = train_clip_frames + context_frames
+
+    def _get_example(self, metadatum):
+        fname, _, _ = metadatum
+        wav_path = self.root.joinpath("wavs", fname + ".wav")
+
+        config = self.config
+        sr = config.sample_rate
+        fft_window_shift = config.fft_window_shift
+        fft_window_size = config.fft_window_size
+        fft_size = config.fft_size
+        
+        audio, loaded_sr = librosa.load(wav_path, sr=None)
+        assert loaded_sr == sr
+
+        # Pad audio to the right size.
+        frames = math.ceil(float(audio.size) / fft_window_shift)
+        fft_padding = (fft_size - fft_window_shift) // 2
+        desired_length = frames * fft_window_shift + fft_padding * 2
+        pad_amount = (desired_length - audio.size) // 2
+        
+        if audio.size % 2 == 0:
+            audio = np.pad(audio, (pad_amount, pad_amount), mode='reflect')
+        else:
+            audio = np.pad(audio, (pad_amount, pad_amount + 1), mode='reflect')
+        
+        # Normalize audio.
+        audio = audio / np.abs(audio).max() * 0.999
+        
+        # Compute mel-spectrogram.
+        # Turn center to False to prevent internal padding.
+        spectrogram = librosa.core.stft(
+            audio, hop_length=fft_window_shift,
+            win_length=fft_window_size, n_fft=fft_size, center=False)
+        spectrogram_magnitude = np.abs(spectrogram)
+        
+        # Compute mel-spectrograms.
+        mel_filter_bank = librosa.filters.mel(sr=sr, n_fft=fft_size,
+                                              n_mels=config.mel_bands)
+        mel_spectrogram = np.dot(mel_filter_bank, spectrogram_magnitude)
+        mel_spectrogram = mel_spectrogram.T
+        
+        # Rescale mel_spectrogram.
+        min_level, ref_level = 1e-5, 20
+        mel_spectrogram = 20 * np.log10(np.maximum(min_level, mel_spectrogram))
+        mel_spectrogram = mel_spectrogram - ref_level
+        mel_spectrogram = np.clip((mel_spectrogram + 100) / 100, 0, 1)
+        
+        # Extract the center of audio that corresponds to mel spectrograms.
+        audio = audio[fft_padding : -fft_padding]
+        assert mel_spectrogram.shape[0] * fft_window_shift == audio.size
+
+        return audio, mel_spectrogram
+
+
+class Subset(dataset.Dataset): 
+    def __init__(self, dataset, indices, valid):
+        self.dataset = dataset
+        self.indices = indices
+        self.valid = valid
+
+    def __getitem__(self, idx):
+        fft_window_shift = self.dataset.fft_window_shift
+        num_frames = self.dataset.num_frames
+        audio, mel = self.dataset[self.indices[idx]]
+
+        if self.valid:
+            audio_start = 0
+        else:
+            # Randomly crop context + train_clip_second of audio.
+            audio_frames = int(audio.size) // fft_window_shift
+            max_start_frame = audio_frames - num_frames
+            assert max_start_frame >= 0, "audio {} is too short".format(idx)
+
+            frame_start = random.randint(0, max_start_frame)
+            frame_end = frame_start + num_frames
+
+            audio_start = frame_start * fft_window_shift
+            audio_end = frame_end * fft_window_shift
+            
+            audio = audio[audio_start : audio_end]
+
+        return audio, mel, audio_start
+
+    def _batch_examples(self, batch):
+        audios = [sample[0] for sample in batch]
+        audio_starts = [sample[2] for sample in batch]
+    
+        # mels shape [num_frames, mel_bands]
+        max_frames = max(sample[1].shape[0] for sample in batch) 
+        mels = [utils.pad_to_size(sample[1], max_frames) for sample in batch]
+        
+        audios = np.array(audios, dtype=np.float32)
+        mels = np.array(mels, dtype=np.float32)
+        audio_starts = np.array(audio_starts, dtype=np.int32)
+    
+        return audios, mels, audio_starts
+
+    def __len__(self):
+        return len(self.indices)
+
+
+class DistributedSampler(Sampler):
+    def __init__(self, dataset_size, num_trainers, rank, shuffle=True):
+        self.dataset_size = dataset_size
+        self.num_trainers = num_trainers
+        self.rank = rank
+        self.num_samples = int(math.ceil(dataset_size / num_trainers))
+        self.total_size = self.num_samples * num_trainers
+        assert self.total_size >= self.dataset_size
+        self.shuffle = shuffle
+    
+    def __iter__(self):
+        indices = list(range(self.dataset_size))
+        if self.shuffle:
+            random.shuffle(indices)
+        
+        # Append extra samples to make it evenly distributed on all trainers.
+        indices += indices[:(self.total_size - self.dataset_size)]
+        assert len(indices) == self.total_size
+
+        # Subset samples for each trainer.
+        indices = indices[self.rank:self.total_size:self.num_trainers]
+        assert len(indices) ==  self.num_samples
+
+        return iter(indices)
+
+    def __len__(self):
+        return self.num_samples
+
+
+class LJSpeech:
+    def __init__(self, config, nranks, rank):
+        place = fluid.CUDAPlace(rank) if config.use_gpu else fluid.CPUPlace()
+
+        # Whole LJSpeech dataset.
+        ds = Dataset(config)
+
+        # Split into train and valid dataset.
+        indices = list(range(len(ds)))
+        train_indices = indices[config.valid_size:]
+        valid_indices = indices[:config.valid_size]
+        random.shuffle(train_indices)
+
+        # Train dataset.
+        trainset = Subset(ds, train_indices, valid=False)
+        sampler = DistributedSampler(len(trainset), nranks, rank) 
+        total_bs = config.batch_size
+        assert total_bs % nranks == 0
+        train_sampler = BatchSampler(sampler, total_bs // nranks,
+            drop_last=True)
+        trainloader = DataCargo(trainset, batch_sampler=train_sampler)
+
+        trainreader = fluid.io.PyReader(capacity=50, return_list=True)
+        trainreader.decorate_batch_generator(trainloader, place)
+        self.trainloader = (data for _ in iter(int, 1)
+            for data in trainreader())
+
+        # Valid dataset.
+        validset = Subset(ds, valid_indices, valid=True)
+        # Currently only support batch_size = 1 for valid loader.
+        validloader = DataCargo(validset, batch_size=1, shuffle=False)
+
+        validreader = fluid.io.PyReader(capacity=20, return_list=True)
+        validreader.decorate_batch_generator(validloader, place) 
+        self.validloader = validreader

parakeet/models/wavenet/ops.py

+import paddle
+from paddle import fluid
+import paddle.fluid.dygraph as dg
+import numpy as np
+
+import weight_norm
+
+
+def Embedding(name_scope,
+              num_embeddings,
+              embed_dim,
+              padding_idx=None,
+              std=0.1,
+              dtype="float32"):
+    # param attrs
+    weight_attr = fluid.ParamAttr(initializer=fluid.initializer.Normal(
+        scale=std))
+    layer = dg.Embedding(
+        name_scope, (num_embeddings, embed_dim),
+        padding_idx=padding_idx,
+        param_attr=weight_attr,
+        dtype=dtype)
+    return layer
+
+
+def FC(name_scope,
+       in_features,
+       size,
+       num_flatten_dims=1,
+       relu=False,
+       dropout=0.0,
+       act=None,
+       dtype="float32"):
+    """
+    A special Linear Layer, when it is used with dropout, the weight is 
+    initialized as normal(0, std=np.sqrt((1-dropout) / in_features))
+    """
+
+    # stds
+    if isinstance(in_features, int):
+        in_features = [in_features]
+
+    stds = [np.sqrt((1.0 - dropout) / in_feature) for in_feature in in_features]
+    if relu:
+        stds = [std * np.sqrt(2.0) for std in stds]
+
+    weight_inits = [
+        fluid.initializer.NormalInitializer(scale=std) for std in stds
+    ]
+    bias_init = fluid.initializer.ConstantInitializer(0.0)
+
+    # param attrs
+    weight_attrs = [fluid.ParamAttr(initializer=init) for init in weight_inits]
+    bias_attr = fluid.ParamAttr(initializer=bias_init)
+
+    layer = weight_norm.FC(name_scope,
+                           size,
+                           num_flatten_dims=num_flatten_dims,
+                           param_attr=weight_attrs,
+                           bias_attr=bias_attr,
+                           act=act,
+                           dtype=dtype)
+    return layer
+
+
+def Conv1D(name_scope,
+           in_channels,
+           num_filters,
+           filter_size=2,
+           dilation=1,
+           groups=None,
+           causal=False,
+           std_mul=1.0,
+           dropout=0.0,
+           use_cudnn=True,
+           act=None,
+           dtype="float32"):
+    """
+    A special Conv1D Layer, when it is used with dropout, the weight is 
+    initialized as 
+    normal(0, std=np.sqrt(std_mul * (1-dropout) / (filter_size * in_channels)))
+    """
+    # std
+    std = np.sqrt((std_mul * (1.0 - dropout)) / (filter_size * in_channels))
+    weight_init = fluid.initializer.NormalInitializer(loc=0.0, scale=std)
+    bias_init = fluid.initializer.ConstantInitializer(0.0)
+
+    # param attrs
+    weight_attr = fluid.ParamAttr(initializer=weight_init)
+    bias_attr = fluid.ParamAttr(initializer=bias_init)
+
+    layer = weight_norm.Conv1D(
+        name_scope,
+        num_filters,
+        filter_size,
+        dilation,
+        groups=groups,
+        causal=causal,
+        param_attr=weight_attr,
+        bias_attr=bias_attr,
+        use_cudnn=use_cudnn,
+        act=act,
+        dtype=dtype)
+    return layer
+
+
+class Conv1D_GU(dg.Layer):
+    def __init__(self,
+                 name_scope,
+                 conditioner_dim,
+                 in_channels,
+                 num_filters,
+                 filter_size,
+                 dilation,
+                 causal=False,
+                 residual=True,
+                 dtype="float32"):
+        super(Conv1D_GU, self).__init__(name_scope, dtype=dtype)
+
+        self.conditioner_dim = conditioner_dim
+        self.in_channels = in_channels
+        self.num_filters = num_filters
+        self.filter_size = filter_size
+        self.dilation = dilation
+        self.causal = causal
+        self.residual = residual
+
+        if residual:
+            assert (
+                in_channels == num_filters
+            ), "this block uses residual connection"\
+                "the input_channels should equals num_filters"
+
+        self.conv = Conv1D(
+            self.full_name(),
+            in_channels,
+            2 * num_filters,
+            filter_size,
+            dilation,
+            causal=causal,
+            dtype=dtype)
+
+        self.fc = Conv1D(
+            self.full_name(),
+            conditioner_dim,
+            2 * num_filters,
+            filter_size=1,
+            dilation=1,
+            causal=False,
+            dtype=dtype)
+
+    def forward(self, x, skip=None, conditioner=None):
+        """
+        Args:
+            x (Variable): Shape(B, C_in, 1, T), the input of Conv1DGLU
+                layer, where B means batch_size, C_in means the input channels
+                T means input time steps.
+            conditioner (Variable): Shape(B, C_con, 1, T), expanded mel
+                conditioner, where C_con is conditioner hidden dim which
+                equals the num of mel bands. Note that when using residual
+                connection, the Conv1DGLU does not change the number of
+                channels, so out channels equals input channels.
+        Returns:
+            x (Variable): Shape(B, C_out, 1, T), the output of Conv1DGLU, where
+                C_out means the output channels of Conv1DGLU.
+        """
+        residual = x
+        x = self.conv(x)
+
+        if conditioner is not None:
+            cond_bias = self.fc(conditioner)
+            x += cond_bias
+
+        content, gate = fluid.layers.split(x, num_or_sections=2, dim=1)
+
+        # Gated Unit.
+        x = fluid.layers.elementwise_mul(fluid.layers.sigmoid(gate),
+                                         fluid.layers.tanh(content))
+
+        if skip is None:
+            skip = x
+        else:
+            skip = fluid.layers.scale(skip + x, np.sqrt(0.5))
+
+        if self.residual:
+            x = fluid.layers.scale(residual + x, np.sqrt(0.5))
+
+        return x, skip
+
+    def add_input(self, x, skip=None, conditioner=None):
+        """
+        Inputs:
+        x: shape(B, num_filters, 1, time_steps)
+        conditioner: shape(B, conditioner_dim, 1, time_steps)
+        Outputs:
+        out: shape(B, num_filters, 1, time_steps), where time_steps = 1
+        """
+        residual = x
+
+        # add step input and produce step output
+        x = self.conv.add_input(x)
+
+        if conditioner is not None:
+            cond_bias = self.fc(conditioner)
+            x += cond_bias
+
+        content, gate = fluid.layers.split(x, num_or_sections=2, dim=1)
+
+        # Gated Unit.
+        x = fluid.layers.elementwise_mul(fluid.layers.sigmoid(gate),
+                                         fluid.layers.tanh(content))
+
+        if skip is None:
+            skip = x
+        else:
+            skip = fluid.layers.scale(skip + x, np.sqrt(0.5))
+
+        if self.residual:
+            x = fluid.layers.scale(residual + x, np.sqrt(0.5))
+
+        return x, skip
+
+
+def Conv2DTranspose(name_scope,
+                    num_filters,
+                    filter_size,
+                    padding=0,
+                    stride=1,
+                    dilation=1,
+                    use_cudnn=True,
+                    act=None,
+                    dtype="float32"):
+    val = 1.0 / (filter_size[0] * filter_size[1])
+    weight_init = fluid.initializer.ConstantInitializer(val)
+    weight_attr = fluid.ParamAttr(initializer=weight_init)
+
+    layer = weight_norm.Conv2DTranspose(
+        name_scope,
+        num_filters,
+        filter_size=filter_size,
+        padding=padding,
+        stride=stride,
+        dilation=dilation,
+        param_attr=weight_attr,
+        use_cudnn=use_cudnn,
+        act=act,
+        dtype=dtype)
+
+    return layer

parakeet/models/wavenet/slurm.py

+"""
+Utility module for restarting training when using SLURM.
+"""
+import subprocess
+import os
+import sys
+import shlex
+import re
+import time
+
+
+def job_info():
+    """Get information about the current job using `scontrol show job`.
+    Returns a dict mapping parameter names (e.g. "UserId", "RunTime", etc) to
+    their values, both as strings.
+    """
+    job_id = int(os.environ["SLURM_JOB_ID"])
+
+    command = ["scontrol", "show", "job", str(job_id)]
+    output = subprocess.check_output(command).decode("utf-8")
+
+    # Use a regex to extract the parameter names and values
+    pattern = "([A-Za-z/]*)=([^ \t\n]*)"
+    return dict(re.findall(pattern, output))
+
+
+def parse_hours(text):
+    """Parse a time format HH or DD-HH into a number of hours."""
+    hour_chunks = text.split("-")
+    if len(hour_chunks) == 1:
+        return int(hour_chunks[0])
+    elif len(hour_chunks) == 2:
+        return 24 * int(hour_chunks[0]) + int(hour_chunks[1])
+    else:
+        raise ValueError("Unexpected hour format (expected HH or "
+                         "DD-HH, but got {}).".format(text))
+
+
+def parse_time(text):
+    """Convert slurm time to an integer.
+    Expects time to be of the form:
+    "hours:minutes:seconds" or "day-hours:minutes:seconds".
+    """
+    hours, minutes, seconds = text.split(":")
+    try:
+        return parse_hours(hours) * 3600 + int(minutes) * 60 + int(seconds)
+    except ValueError as e:
+        raise ValueError("Error parsing time {}. Got error {}.".format(
+            text, str(e)))
+
+
+def restart_command():
+    """Using the environment and SLURM command, create a command that, when,
+    run, will enqueue a repeat of the current job using `sbatch`.
+    Return the command as a list of strings, suitable for passing to
+    `subprocess.check_call` or similar functions.
+    Returns:
+        resume_command: list<str>, command to run to restart job.
+        end_time: int or None; the time the job will end or None
+            if the job has unlimited runtime.
+    """
+    # Make sure `RunTime` could be parsed correctly.
+    while job_info()["RunTime"] == "INVALID":
+        time.sleep(1)
+
+    # Get all the necessary information by querying SLURM with this job id
+    info = job_info()
+
+    try:
+        num_cpus = int(info["CPUs/Task"])
+    except KeyError:
+        num_cpus = int(os.environ["SLURM_CPUS_PER_TASK"])
+
+    num_tasks = int(os.environ["SLURM_NTASKS"])
+    nodes = info["NumNodes"]
+    gres, partition = info.get("Gres"), info.get("Partition")
+    stderr, stdout = info.get("StdErr"), info.get("StdOut")
+    job_name = info.get("JobName")
+    command = ["sbatch", "--job-name={}".format(job_name),
+               "--ntasks={}".format(num_tasks)]
+
+    if partition:
+        command.extend(["--partition", partition])
+
+    if gres and gres != "(null)":
+        command.extend(["--gres", gres])
+        num_gpu = int(gres.split(':')[-1])
+        print("number of gpu assigned by slurm is {}".format(num_gpu))
+
+    if stderr:
+        command.extend(["--error", stderr])
+
+    if stdout:
+        command.extend(["--output", stdout])
+
+    python = subprocess.check_output(
+        ["/usr/bin/which", "python3"]).decode("utf-8").strip()
+    dist_setting = ['-m', 'paddle.distributed.launch']
+    wrap_cmd = ["srun", python, '-u'] + dist_setting + sys.argv
+
+    command.append(
+        "--wrap={}".format(" ".join(shlex.quote(arg) for arg in wrap_cmd)))
+    time_limit_string = info["TimeLimit"]
+    if time_limit_string.lower() == "unlimited":
+        print("UNLIMITED detected: restart OFF, infinite learning ON.",
+              flush=True)
+        return command, None
+    time_limit = parse_time(time_limit_string)
+    runtime = parse_time(info["RunTime"])
+    end_time = time.time() + time_limit - runtime
+
+    return command, end_time

parakeet/models/wavenet/synthesis.py

+import os
+import random
+from pprint import pprint
+
+import jsonargparse
+import numpy as np
+import paddle.fluid.dygraph as dg
+from paddle import fluid
+
+import utils
+from wavenet import WaveNet
+
+
+def add_options_to_parser(parser):
+    parser.add_argument('--model', type=str, default='wavenet',
+        help="general name of the model")
+    parser.add_argument('--name', type=str,
+        help="specific name of the training model")
+    parser.add_argument('--root', type=str,
+        help="root path of the LJSpeech dataset")
+
+    parser.add_argument('--use_gpu', type=bool, default=True,
+        help="option to use gpu training")
+
+    parser.add_argument('--iteration', type=int, default=None,
+        help=("which iteration of checkpoint to load, "
+              "default to load the latest checkpoint"))
+    parser.add_argument('--checkpoint', type=str, default=None,
+        help="path of the checkpoint to load")
+
+    parser.add_argument('--output', type=str, default="./syn_audios",
+        help="path to write synthesized audio files")
+    parser.add_argument('--sample', type=int,
+        help="which of the valid samples to synthesize audio")
+
+
+def synthesize(config):
+    pprint(jsonargparse.namespace_to_dict(config))
+
+    # Get checkpoint directory path.
+    run_dir = os.path.join("runs", config.model, config.name)
+    checkpoint_dir = os.path.join(run_dir, "checkpoint")
+
+    # Configurate device.
+    place = fluid.CUDAPlace(0) if config.use_gpu else fluid.CPUPlace()
+
+    with dg.guard(place):
+        # Fix random seed.
+        seed = config.seed
+        random.seed(seed)
+        np.random.seed(seed)
+        fluid.default_startup_program().random_seed = seed
+        fluid.default_main_program().random_seed = seed
+        print("Random Seed: ", seed)
+        
+        # Build model.
+        model = WaveNet(config, checkpoint_dir)
+        model.build(training=False)
+
+        # Obtain the current iteration.
+        if config.checkpoint is None:
+            if config.iteration is None:
+                iteration = utils.load_latest_checkpoint(checkpoint_dir)
+            else:
+                iteration = config.iteration
+        else:
+            iteration = int(config.checkpoint.split('/')[-1].split('-')[-1])
+
+        # Run model inference.
+        model.infer(iteration)
+
+
+if __name__ == "__main__":
+    # Create parser.
+    parser = jsonargparse.ArgumentParser(
+        description="Synthesize audio using WaveNet model",
+        formatter_class='default_argparse')
+    add_options_to_parser(parser)
+    utils.add_config_options_to_parser(parser)
+
+    # Parse argument from both command line and yaml config file.
+    # For conflicting updates to the same field,
+    # the preceding update will be overwritten by the following one.
+    config = parser.parse_args()
+    synthesize(config)

parakeet/models/wavenet/train.py

+import os
+import random
+import subprocess
+import time
+from pprint import pprint
+
+import jsonargparse
+import numpy as np
+import paddle.fluid.dygraph as dg
+from paddle import fluid
+from tensorboardX import SummaryWriter
+
+import slurm
+import utils
+from wavenet import WaveNet
+
+MAXIMUM_SAVE_TIME = 10 * 60
+
+
+def add_options_to_parser(parser):
+    parser.add_argument('--model', type=str, default='wavenet',
+        help="general name of the model")
+    parser.add_argument('--name', type=str,
+        help="specific name of the training model")
+    parser.add_argument('--root', type=str,
+        help="root path of the LJSpeech dataset")
+
+    parser.add_argument('--parallel', type=bool, default=True,
+        help="option to use data parallel training")
+    parser.add_argument('--use_gpu', type=bool, default=True,
+        help="option to use gpu training")
+
+    parser.add_argument('--iteration', type=int, default=None,
+        help=("which iteration of checkpoint to load, "
+              "default to load the latest checkpoint"))
+    parser.add_argument('--checkpoint', type=str, default=None,
+        help="path of the checkpoint to load")
+    parser.add_argument('--slurm', type=bool, default=False,
+        help="whether you are using slurm to submit training jobs")
+
+
+def train(config):
+    use_gpu = config.use_gpu
+    parallel = config.parallel if use_gpu else False
+
+    # Get the rank of the current training process.
+    rank = dg.parallel.Env().local_rank if parallel else 0
+    nranks = dg.parallel.Env().nranks if parallel else 1
+
+    if rank == 0:
+        # Print the whole config setting.
+        pprint(jsonargparse.namespace_to_dict(config))
+
+    # Make checkpoint directory.
+    run_dir = os.path.join("runs", config.model, config.name)
+    checkpoint_dir = os.path.join(run_dir, "checkpoint")
+    os.makedirs(checkpoint_dir, exist_ok=True)
+
+    # Create tensorboard logger.
+    tb = SummaryWriter(os.path.join(run_dir, "logs")) \
+        if rank == 0 else None
+
+    # Configurate device
+    place = fluid.CUDAPlace(rank) if use_gpu else fluid.CPUPlace()
+
+    with dg.guard(place):
+        # Fix random seed.
+        seed = config.seed
+        random.seed(seed)
+        np.random.seed(seed)
+        fluid.default_startup_program().random_seed = seed
+        fluid.default_main_program().random_seed = seed
+        print("Random Seed: ", seed)
+
+        # Build model.
+        model = WaveNet(config, checkpoint_dir, parallel, rank, nranks, tb)
+        model.build()
+
+        # Obtain the current iteration.
+        if config.checkpoint is None:
+            if config.iteration is None:
+                iteration = utils.load_latest_checkpoint(checkpoint_dir, rank)
+            else:
+                iteration = config.iteration
+        else:
+            iteration = int(config.checkpoint.split('/')[-1].split('-')[-1])
+
+        # Get restart command if using slurm.
+        if config.slurm:
+            resume_command, death_time = slurm.restart_command()
+            if rank == 0:
+                print("Restart command:", " ".join(resume_command))
+        done = False
+
+        while iteration < config.max_iterations:
+            # Run one single training step.
+            model.train_step(iteration)
+
+            iteration += 1
+
+            if iteration % config.test_every == 0:
+                # Run validation step.
+                model.valid_step(iteration)
+
+            # Check whether reaching the time limit.
+            if config.slurm:
+                done = (death_time is not None and death_time - time.time() <
+                    MAXIMUM_SAVE_TIME)
+
+            if rank == 0 and done:
+                print("Saving progress before exiting.")
+                model.save(iteration)
+
+                print("Running restart command:", " ".join(resume_command))
+                # Submit restart command.
+                subprocess.check_call(resume_command)
+                break
+
+            if rank == 0 and iteration % config.save_every == 0:
+                # Save parameters.
+                model.save(iteration)
+
+    # Close TensorBoard.
+    if rank == 0:
+        tb.close()
+
+
+if __name__ == "__main__":
+    # Create parser.
+    parser = jsonargparse.ArgumentParser(description="Train WaveNet model",
+        formatter_class='default_argparse')
+    add_options_to_parser(parser)
+    utils.add_config_options_to_parser(parser)
+
+    # Parse argument from both command line and yaml config file.
+    # For conflicting updates to the same field, 
+    # the preceding update will be overwritten by the following one.
+    config = parser.parse_args()
+    train(config) 

parakeet/models/wavenet/utils.py

+import itertools
+import os
+import time
+
+import jsonargparse
+import numpy as np
+import paddle.fluid.dygraph as dg
+
+
+def add_config_options_to_parser(parser):
+    parser.add_argument('--valid_size', type=int,
+        help="size of the valid dataset")
+    parser.add_argument('--train_clip_second', type=float,
+        help="the length of audio clip for training")
+    parser.add_argument('--sample_rate', type=int,
+        help="sampling rate of audio data file")
+    parser.add_argument('--fft_window_shift', type=int,
+        help="the shift of fft window for each frame")
+    parser.add_argument('--fft_window_size', type=int,
+        help="the size of fft window for each frame")
+    parser.add_argument('--fft_size', type=int,
+        help="the size of fft filter on each frame")
+    parser.add_argument('--mel_bands', type=int,
+        help="the number of mel bands when calculating mel spectrograms")
+
+    parser.add_argument('--seed', type=int,
+        help="seed of random initialization for the model")
+    parser.add_argument('--batch_size', type=int,
+        help="batch size for training")
+    parser.add_argument('--test_every', type=int,
+        help="test interval during training")
+    parser.add_argument('--save_every', type=int,
+        help="checkpointing interval during training")
+    parser.add_argument('--max_iterations', type=int,
+        help="maximum training iterations")
+
+    parser.add_argument('--layers', type=int,
+        help="number of dilated convolution layers")
+    parser.add_argument('--kernel_width', type=int,
+        help="dilated convolution kernel width")
+    parser.add_argument('--dilation_block', type=list,
+        help="dilated convolution kernel width")
+    parser.add_argument('--residual_channels', type=int)
+    parser.add_argument('--skip_channels', type=int)
+    parser.add_argument('--loss_type', type=str,
+        help="mix-gaussian-pdf or softmax")
+    parser.add_argument('--num_channels', type=int, default=None,
+        help="number of channels for softmax output")
+    parser.add_argument('--num_mixtures', type=int, default=None,
+        help="number of gaussian mixtures for gaussian output")
+    parser.add_argument('--log_scale_min', type=float, default=None,
+        help="minimum clip value of log variance of gaussian output")
+
+    parser.add_argument('--conditioner.filter_sizes', type=list,
+        help="conv2d tranpose op filter sizes for building conditioner")
+    parser.add_argument('--conditioner.upsample_factors', type=list,
+        help="list of upsample factors for building conditioner")
+
+    parser.add_argument('--learning_rate', type=float)
+    parser.add_argument('--gradient_max_norm', type=float)
+    parser.add_argument('--anneal.every', type=int,
+        help="step interval for annealing learning rate")
+    parser.add_argument('--anneal.rate', type=float)
+
+    parser.add_argument('--config', action=jsonargparse.ActionConfigFile)
+
+
+def pad_to_size(array, length, pad_with=0.0):
+    """
+    Pad an array on the first (length) axis to a given length.
+    """
+    padding = length - array.shape[0]
+    assert padding >= 0, "Padding required was less than zero"
+
+    paddings = [(0, 0)] * len(array.shape)
+    paddings[0] = (0, padding)
+
+    return np.pad(array, paddings, mode='constant', constant_values=pad_with)
+
+
+def calculate_context_size(config):
+    dilations = list(
+        itertools.islice(
+            itertools.cycle(config.dilation_block), config.layers))
+    config.context_size = sum(dilations) + 1
+    print("Context size is", config.context_size)
+
+
+def load_latest_checkpoint(checkpoint_dir, rank=0):
+    checkpoint_path = os.path.join(checkpoint_dir, "checkpoint")
+    # Create checkpoint index file if not exist.
+    if (not os.path.isfile(checkpoint_path)) and rank == 0:
+        with open(checkpoint_path, "w") as handle:
+            handle.write("model_checkpoint_path: step-0")
+
+    # Make sure that other process waits until checkpoint file is created
+    # by process 0.
+    while not os.path.isfile(checkpoint_path):
+        time.sleep(1)
+
+    # Fetch the latest checkpoint index.
+    with open(checkpoint_path, "r") as handle:
+        latest_checkpoint = handle.readline().split()[-1]
+        iteration = int(latest_checkpoint.split("-")[-1])
+
+    return iteration
+
+
+def save_latest_checkpoint(checkpoint_dir, iteration):
+    checkpoint_path = os.path.join(checkpoint_dir, "checkpoint")
+    # Update the latest checkpoint index.
+    with open(checkpoint_path, "w") as handle:
+        handle.write("model_checkpoint_path: step-{}".format(iteration))
+
+
+def load_parameters(checkpoint_dir, rank, model, optimizer=None,
+                    iteration=None, file_path=None):
+    if file_path is None:
+        if iteration is None:
+            iteration = load_latest_checkpoint(checkpoint_dir, rank)
+        if iteration == 0:
+            return
+        file_path = "{}/step-{}".format(checkpoint_dir, iteration)
+
+    model_dict, optimizer_dict = dg.load_dygraph(file_path)
+    model.set_dict(model_dict)
+    print("[checkpoint] Rank {}: loaded model from {}".format(rank, file_path))
+    if optimizer and optimizer_dict:
+        optimizer.set_dict(optimizer_dict)
+        print("[checkpoint] Rank {}: loaded optimizer state from {}".format(
+              rank, file_path))
+
+
+def save_latest_parameters(checkpoint_dir, iteration, model, optimizer=None):
+    file_path = "{}/step-{}".format(checkpoint_dir, iteration)
+    model_dict = model.state_dict()
+    dg.save_dygraph(model_dict, file_path)
+    print("[checkpoint] Saved model to {}".format(file_path))
+
+    if optimizer:
+        opt_dict = optimizer.state_dict()
+        dg.save_dygraph(opt_dict, file_path)
+        print("[checkpoint] Saved optimzier state to {}".format(file_path))

parakeet/models/wavenet/wavenet.py

+import itertools
+import os
+import time
+
+import librosa
+import numpy as np
+from paddle import fluid
+import paddle.fluid.dygraph as dg
+
+import utils
+from data import LJSpeech
+from wavenet_modules import WaveNetModule, debug
+
+
+class WaveNet():
+    def __init__(self, config, checkpoint_dir, parallel=False, rank=0,
+                 nranks=1, tb_logger=None):
+        # Process config to calculate the context size
+        dilations = list(
+            itertools.islice(
+                itertools.cycle(config.dilation_block), config.layers))
+        config.context_size = sum(dilations) + 1
+        self.config = config
+        self.checkpoint_dir = checkpoint_dir
+        self.parallel = parallel
+        self.rank = rank
+        self.nranks = nranks
+        self.tb_logger = tb_logger
+
+    def build(self, training=True):
+        config = self.config
+        dataset = LJSpeech(config, self.nranks, self.rank) 
+        self.trainloader = dataset.trainloader
+        self.validloader = dataset.validloader
+
+#        if self.rank == 0:
+#            for i, (audios, mels, ids) in enumerate(self.validloader()):
+#                print("audios {}, mels {}, ids {}".format(audios.dtype, mels.dtype, ids.dtype))
+#                print("{}: rank {}, audios {}, mels {}, indices {} / {}".format(
+#                    i, self.rank, audios.shape, mels.shape, ids.shape,
+#                    ids.numpy()))
+#    
+#            for i, (audios, mels, ids) in enumerate(self.trainloader):
+#                print("{}: rank {}, audios {}, mels {}, indices {} / {}".format(
+#                    i, self.rank, audios.shape, mels.shape, ids.shape,
+#                    ids.numpy()))
+
+        wavenet = WaveNetModule("wavenet", config, self.rank)
+        
+        # Dry run once to create and initalize all necessary parameters.
+        audio = dg.to_variable(np.random.randn(1, 20000).astype(np.float32))
+        mel = dg.to_variable(
+            np.random.randn(1, 100, self.config.mel_bands).astype(np.float32))
+        audio_start = dg.to_variable(np.array([0], dtype=np.int32))
+        wavenet(audio, mel, audio_start)
+
+        if training:
+            # Create Learning rate scheduler.
+            lr_scheduler = dg.ExponentialDecay(
+                learning_rate = config.learning_rate,
+                decay_steps = config.anneal.every,
+                decay_rate = config.anneal.rate,
+                staircase=True)
+    
+            optimizer = fluid.optimizer.AdamOptimizer(
+                learning_rate=lr_scheduler)
+    
+            clipper = fluid.dygraph_grad_clip.GradClipByGlobalNorm(
+                config.gradient_max_norm)
+
+            # Load parameters.
+            utils.load_parameters(self.checkpoint_dir, self.rank,
+                                  wavenet, optimizer,
+                                  iteration=config.iteration,
+                                  file_path=config.checkpoint)
+            print("Rank {}: checkpoint loaded.".format(self.rank))
+    
+            # Data parallelism.
+            if self.parallel:
+                strategy = dg.parallel.prepare_context()
+                wavenet = dg.parallel.DataParallel(wavenet, strategy)
+    
+            self.wavenet = wavenet
+            self.optimizer = optimizer
+            self.clipper = clipper
+
+        else:
+            # Load parameters.
+            utils.load_parameters(self.checkpoint_dir, self.rank, wavenet,
+                                  iteration=config.iteration,
+                                  file_path=config.checkpoint)
+            print("Rank {}: checkpoint loaded.".format(self.rank))
+
+            self.wavenet = wavenet
+
+    def train_step(self, iteration):
+        self.wavenet.train()
+
+        start_time = time.time()
+        audios, mels, audio_starts = next(self.trainloader)
+        load_time = time.time()
+
+        loss, _ = self.wavenet(audios, mels, audio_starts)
+
+        if self.parallel:
+            # loss = loss / num_trainers
+            loss = self.wavenet.scale_loss(loss)
+            loss.backward()
+            self.wavenet.apply_collective_grads()
+        else:
+            loss.backward()
+
+        if isinstance(self.optimizer._learning_rate,
+                      fluid.optimizer.LearningRateDecay):
+            current_lr = self.optimizer._learning_rate.step().numpy()
+        else:
+            current_lr = self.optimizer._learning_rate
+
+        self.optimizer.minimize(loss, grad_clip=self.clipper,
+            parameter_list=self.wavenet.parameters())
+        self.wavenet.clear_gradients()
+
+        graph_time = time.time()
+
+        if self.rank == 0:
+            loss_val = float(loss.numpy()) * self.nranks
+            log = "Rank: {} Step: {:^8d} Loss: {:<8.3f} " \
+                  "Time: {:.3f}/{:.3f}".format(
+                  self.rank, iteration, loss_val,
+                  load_time - start_time, graph_time - load_time)
+            print(log)
+
+            tb = self.tb_logger
+            tb.add_scalar("Train-Loss-Rank-0", loss_val, iteration)
+            tb.add_scalar("Learning-Rate", current_lr, iteration)
+
+    @dg.no_grad
+    def valid_step(self, iteration):
+        self.wavenet.eval()
+
+        total_loss = []
+        start_time = time.time()
+        sample_audios = []
+        for audios, mels, audio_starts in self.validloader():
+            loss, sample_audio = self.wavenet(audios, mels, audio_starts, True)
+            total_loss.append(float(loss.numpy()))
+            sample_audios.append(sample_audio)
+        total_time = time.time() - start_time
+
+        if self.rank == 0:
+            loss_val = np.mean(total_loss)
+            log = "Test | Rank: {} AvgLoss: {:<8.3f} Time {:<8.3f}".format(
+                self.rank, loss_val, total_time)
+            print(log)
+
+            tb = self.tb_logger
+            tb.add_scalar("Valid-Avg-Loss", loss_val, iteration)
+            tb.add_audio("Teacher-Forced-Audio-0", sample_audios[0].numpy(),
+                iteration, sample_rate=self.config.sample_rate)
+            tb.add_audio("Teacher-Forced-Audio-1", sample_audios[1].numpy(),
+                iteration, sample_rate=self.config.sample_rate)
+
+    def save(self, iteration):
+        utils.save_latest_parameters(self.checkpoint_dir, iteration,
+                                     self.wavenet, self.optimizer)
+        utils.save_latest_checkpoint(self.checkpoint_dir, iteration)
+
+    @dg.no_grad
+    def infer(self, iteration):
+        self.wavenet.eval()
+
+        config = self.config
+        sample = config.sample
+
+        output = "{}/{}/iter-{}".format(config.output, config.name, iteration)
+        os.makedirs(output, exist_ok=True)
+
+        filename = "{}/valid_{}.wav".format(output, sample)
+        print("Synthesize sample {}, save as {}".format(sample, filename))
+
+        mels_list = [mels for _, mels, _ in self.validloader()]
+        start_time = time.time()
+        syn_audio = self.wavenet.synthesize(mels_list[sample])
+        syn_time = time.time() - start_time
+        print("audio shape {}, synthesis time {}".format(
+            syn_audio.shape, syn_time))
+        librosa.output.write_wav(filename, syn_audio,
+            sr=config.sample_rate)

parakeet/models/wavenet/wavenet_modules.py

+import itertools
+import math
+
+import numpy as np
+from paddle import fluid
+import paddle.fluid.dygraph as dg
+import ops
+import weight_norm
+
+
+def get_padding(filter_size, stride, padding_type='same'):
+    if padding_type == 'same':
+        padding = [(x - y) // 2 for x, y in zip(filter_size, stride)]
+    else:
+        raise ValueError("Only support same padding")
+    return padding
+
+
+def debug(x, var_name, rank, verbose=False):
+    if not verbose and rank != 0:
+        return
+    dim = len(x.shape)
+    if not isinstance(x, np.ndarray):
+        x = x.numpy()
+    if dim == 1:
+        print("Rank {}".format(rank), var_name, "shape {}, value {}".format(x.shape, x))
+    elif dim == 2:
+        print("Rank {}".format(rank), var_name, "shape {}, value {}".format(x.shape, x[:, :5]))
+    elif dim == 3:
+        print("Rank {}".format(rank), var_name, "shape {}, value {}".format(x.shape, x[:, :5, 0]))
+    else:
+        print("Rank", rank, var_name, "shape", x.shape)
+
+
+def extract_slices(x, audio_starts, audio_length, rank):
+    slices = []
+    for i in range(x.shape[0]): 
+        start = audio_starts.numpy()[i]
+        end = start + audio_length
+        slice = fluid.layers.slice(
+            x, axes=[0, 1], starts=[i, start], ends=[i+1, end])
+        slices.append(fluid.layers.squeeze(slice, [0]))
+
+    x = fluid.layers.stack(slices, axis=0)
+
+    return x
+
+
+class Conditioner(dg.Layer):
+    def __init__(self, name_scope, config):
+        super(Conditioner, self).__init__(name_scope)
+        upsample_factors = config.conditioner.upsample_factors
+        filter_sizes = config.conditioner.filter_sizes
+        assert np.prod(upsample_factors) == config.fft_window_shift
+
+        self.deconvs = []
+        for i, up_scale in enumerate(upsample_factors):
+            stride = (up_scale, 1)
+            padding = get_padding(filter_sizes[i], stride)
+            self.deconvs.append(
+                ops.Conv2DTranspose(
+                    self.full_name(),
+                    num_filters=1,
+                    filter_size=filter_sizes[i],
+                    padding=padding,
+                    stride=stride))
+
+        # Register python list as parameters.
+        for i, layer in enumerate(self.deconvs):
+            self.add_sublayer("conv_transpose_{}".format(i), layer)
+        
+    def forward(self, x):
+        x = fluid.layers.unsqueeze(x, 1)
+        for layer in self.deconvs:
+            x = fluid.layers.leaky_relu(layer(x), alpha=0.4)
+
+        return fluid.layers.squeeze(x, [1])
+
+
+class WaveNetModule(dg.Layer):
+    def __init__(self, name_scope, config, rank):
+        super(WaveNetModule, self).__init__(name_scope)
+        
+        self.rank = rank
+        self.conditioner = Conditioner(self.full_name(), config)
+        self.dilations = list(
+            itertools.islice(
+                itertools.cycle(config.dilation_block), config.layers))
+        self.context_size = sum(self.dilations) + 1
+        self.log_scale_min = config.log_scale_min
+        self.config = config
+
+        print("dilations", self.dilations)
+        print("context_size", self.context_size)
+
+        if config.loss_type == "softmax":
+            self.embedding_fc = ops.Embedding(
+                self.full_name(),
+                num_embeddings=config.num_channels,
+                embed_dim=config.residual_channels)
+        elif config.loss_type == "mix-gaussian-pdf":
+            self.embedding_fc = ops.FC(
+                self.full_name(),
+                in_features=1,
+                size=config.residual_channels,
+                num_flatten_dims=2,
+                relu=False)
+        else:
+            raise ValueError(
+                "loss_type {} is unsupported!".format(loss_type))
+
+        self.dilated_causal_convs = []
+        for dilation in self.dilations:
+            self.dilated_causal_convs.append(
+                ops.Conv1D_GU(
+                    self.full_name(),
+                    conditioner_dim=config.mel_bands,
+                    in_channels=config.residual_channels,
+                    num_filters=config.residual_channels,
+                    filter_size=config.kernel_width,
+                    dilation=dilation,
+                    causal=True
+                )
+            )
+
+        for i, layer in enumerate(self.dilated_causal_convs):
+            self.add_sublayer("dilated_causal_conv_{}".format(i), layer) 
+
+        self.fc1 = ops.FC(
+            self.full_name(),
+            in_features=config.residual_channels,
+            size=config.skip_channels,
+            num_flatten_dims=2,
+            relu=True,
+            act="relu")
+
+        self.fc2 = ops.FC(
+            self.full_name(),
+            in_features=config.skip_channels,
+            size=config.skip_channels,
+            num_flatten_dims=2,
+            relu=True,
+            act="relu")
+
+        if config.loss_type == "softmax":
+            self.fc3 = ops.FC(
+                self.full_name(),
+                in_features=config.skip_channels,
+                size=config.num_channels,
+                num_flatten_dims=2,
+                relu=False)
+        elif config.loss_type == "mix-gaussian-pdf":
+            self.fc3 = ops.FC(
+                self.full_name(),
+                in_features=config.skip_channels,
+                size=3 * config.num_mixtures,
+                num_flatten_dims=2,
+                relu=False)
+        else:
+            raise ValueError(
+                "loss_type {} is unsupported!".format(loss_type))
+
+    def sample_softmax(self, mix_parameters):
+        batch, length, hidden = mix_parameters.shape
+        mix_param_2d = fluid.layers.reshape(mix_parameters,
+            [batch * length, hidden])
+        mix_param_2d = fluid.layers.softmax(mix_param_2d, axis=-1)
+
+        # quantized: [batch * length]
+        quantized = fluid.layers.cast(fluid.layers.sampling_id(mix_param_2d),
+            dtype="float32")
+        samples = (quantized + 0.5) * (2.0 / self.config.num_channels) - 1.0
+
+        # samples: [batch * length]
+        return samples
+
+    def sample_mix_gaussian(self, mix_parameters):
+        # mix_parameters reshape from [bs, 13799, 3 * num_mixtures]
+        # to [bs * 13799, 3 * num_mixtures].
+        batch, length, hidden = mix_parameters.shape
+        mix_param_2d = fluid.layers.reshape(mix_parameters,
+            [batch * length, hidden])
+        K = hidden // 3
+
+        # Unpack the parameters of the mixture of gaussian.
+        logits_pi = mix_param_2d[:, 0 : K]
+        mu = mix_param_2d[:, K : 2*K]
+        log_s = mix_param_2d[:, 2*K : 3*K]
+        s = fluid.layers.exp(log_s)
+
+        pi = fluid.layers.softmax(logits_pi, axis=-1)
+        comp_samples = fluid.layers.sampling_id(pi)
+        
+        row_idx = dg.to_variable(np.arange(batch * length))
+        comp_samples = fluid.layers.stack([row_idx, comp_samples], axis=-1)
+
+        mu_comp = fluid.layers.gather_nd(mu, comp_samples)
+        s_comp = fluid.layers.gather_nd(s, comp_samples) 
+
+        # N(0, 1) Normal Sample.
+        u = fluid.layers.gaussian_random(shape=[batch * length])
+        samples = mu_comp + u * s_comp
+        samples = fluid.layers.clip(samples, min=-1.0, max=1.0)
+
+        return samples
+
+    def softmax_loss(self, targets, mix_parameters):
+        # targets: [bs, 13799] -> [bs, 11752]
+        # mix_params: [bs, 13799, 3] -> [bs, 11752, 3]
+        targets = targets[:, self.context_size:]
+        mix_parameters = mix_parameters[:, self.context_size:, :]
+
+        # Quantized audios to integral values with range [0, num_channels)
+        num_channels = self.config.num_channels
+        targets = fluid.layers.clip(targets, min=-1.0, max=0.99999)
+        quantized = fluid.layers.cast(
+            (targets + 1.0) / 2.0 * num_channels, dtype="int64")
+
+        # per_sample_loss shape: [bs, 17952, 1]
+        per_sample_loss = fluid.layers.softmax_with_cross_entropy(
+            logits=mix_parameters, label=fluid.layers.unsqueeze(quantized, 2))
+        loss = fluid.layers.reduce_mean(per_sample_loss)
+        #debug(loss, "softmax loss", self.rank)
+
+        return loss
+
+    def mixture_density_loss(self, targets, mix_parameters, log_scale_min):
+        # targets: [bs, 13799] -> [bs, 11752]
+        # mix_params: [bs, 13799, 3] -> [bs, 11752, 3]
+        targets = targets[:, self.context_size:]
+        mix_parameters = mix_parameters[:, self.context_size:, :]
+
+        # log_s: [bs, 11752, num_mixture]
+        logits_pi, mu, log_s = fluid.layers.split(mix_parameters, num_or_sections=3, dim=-1)
+
+        pi = fluid.layers.softmax(logits_pi, axis=-1)
+        log_s = fluid.layers.clip(log_s, min=log_scale_min, max=100.0)
+        inv_s = fluid.layers.exp(0.0 - log_s)
+
+        # Calculate gaussian loss.
+        targets = fluid.layers.unsqueeze(targets, -1)
+        targets = fluid.layers.expand(targets, [1, 1, self.config.num_mixtures])
+        x_std =  inv_s * (targets - mu)
+        exponent = fluid.layers.exp(-0.5 * x_std * x_std)
+        # pdf_x: [bs, 11752, 1]
+        pdf_x = 1.0 / np.sqrt(2.0 * np.pi) * inv_s * exponent
+        pdf_x = pi * pdf_x
+        # pdf_x: [bs, 11752]
+        pdf_x = fluid.layers.reduce_sum(pdf_x, dim=-1)
+        per_sample_loss = 0.0 - fluid.layers.log(pdf_x + 1e-9)
+
+        loss = fluid.layers.reduce_mean(per_sample_loss)
+
+        return loss
+
+    def forward(self, audios, mels, audio_starts, sample=False):
+        # audios: [bs, 13800], mels: [bs, full_frame_length, 80]
+        # audio_starts: [bs]
+        # Build conditioner based on mels.
+        full_conditioner = self.conditioner(mels)
+
+        # Slice conditioners.
+        audio_length = audios.shape[1]
+        conditioner = extract_slices(full_conditioner,
+            audio_starts, audio_length, self.rank)
+    
+        # input_audio, target_audio: [bs, 13799]
+        input_audios = audios[:, :-1]
+        target_audios = audios[:, 1:]
+        # conditioner: [bs, 13799, 80]
+        conditioner = conditioner[:, 1:, :]
+
+        loss_type = self.config.loss_type
+
+        # layer_input: [bs, 13799, 128]
+        if loss_type == "softmax":
+            input_audios = fluid.layers.clip(
+                input_audios, min=-1.0, max=0.99999)
+            # quantized have values in [0, num_channels)
+            quantized = fluid.layers.cast(
+                (input_audios + 1.0) / 2.0 * self.config.num_channels,
+                dtype="int64")
+            layer_input = self.embedding_fc(fluid.layers.unsqueeze(quantized, 2))
+        elif loss_type == "mix-gaussian-pdf":
+            layer_input = self.embedding_fc(fluid.layers.unsqueeze(input_audios, 2))
+        else:
+            raise ValueError(
+                "loss_type {} is unsupported!".format(loss_type))
+
+        # layer_input: [bs, res_channel, 1, 13799]
+        layer_input = fluid.layers.unsqueeze(fluid.layers.transpose(layer_input, perm=[0, 2, 1]), 2)
+        # conditioner: [bs, mel_bands, 1, 13799]
+        conditioner = fluid.layers.unsqueeze(fluid.layers.transpose(conditioner, perm=[0, 2, 1]), 2)
+
+        # layer_input: [bs, res_channel, 1, 13799]
+        # skip: [bs, res_channel, 1, 13799]
+        skip = None
+        for i, layer in enumerate(self.dilated_causal_convs):
+            layer_input, skip = layer(layer_input, skip, conditioner)
+            #debug(layer_input, "layer_input_" + str(i), self.rank)
+            #debug(skip, "skip_" + str(i), self.rank)
+
+        # Reshape skip to [bs, 13799, res_channel]
+        skip = fluid.layers.transpose(fluid.layers.squeeze(skip, [2]), perm=[0, 2, 1])
+        #debug(skip, "skip", self.rank)
+
+        # mix_param: [bs, 13799, 3 * num_mixtures]
+        mix_parameters = self.fc3(self.fc2(self.fc1(skip)))
+
+        # Sample teacher-forced audio.
+        sample_audios = None
+        if sample:
+            if loss_type == "softmax":
+                sample_audios = self.sample_softmax(mix_parameters)
+            elif loss_type == "mix-gaussian-pdf":
+                sample_audios = self.sample_mix_gaussian(mix_parameters)
+            else:
+                raise ValueError(
+                    "loss_type {} is unsupported!".format(loss_type))
+            #debug(sample_audios, "sample_audios", self.rank)
+
+        # Calculate mix-gaussian density loss.
+        # padding is all zero.
+        # target_audio: [bs, 13799].
+        # mix_params: [bs, 13799, 3].
+        if loss_type == "softmax":
+            loss = self.softmax_loss(target_audios, mix_parameters)
+        elif loss_type == "mix-gaussian-pdf":
+            loss = self.mixture_density_loss(target_audios,
+                mix_parameters, self.log_scale_min)
+        else:
+            raise ValueError(
+                "loss_type {} is unsupported!".format(loss_type))
+
+        #print("Rank {}, loss {}".format(self.rank, loss.numpy()))
+ 
+        return loss, sample_audios
+
+    def synthesize(self, mels):
+        self.start_new_sequence()
+        print("input mels shape", mels.shape)
+        # mels: [bs=1, n_frames, 80]
+        # conditioner: [1, n_frames * samples_per_frame, 80]
+        # Should I move forward by one sample? No difference
+        # Append context frame to mels
+        bs, n_frames, mel_bands = mels.shape 
+        #num_pad_frames = int(np.ceil(self.context_size / self.config.fft_window_shift))
+        #silence = fluid.layers.zeros(shape=[bs, num_pad_frames, mel_bands], dtype="float32")
+        #inf_mels = fluid.layers.concat([silence, mels], axis=1)
+        #print("padded mels shape", inf_mels.shape)
+
+        #conditioner = self.conditioner(inf_mels)[:, self.context_size:, :]
+        conditioner = self.conditioner(mels)
+        time_steps = conditioner.shape[1]
+        print("Total steps", time_steps)
+
+        loss_type = self.config.loss_type
+        audio_samples = []
+        current_sample = fluid.layers.zeros(shape=[bs, 1, 1], dtype="float32")
+        for i in range(time_steps):
+            if i % 100 == 0:
+                print("Step", i)
+
+            # convert from real value sample to audio embedding.
+            # [bs, 1, 128]
+            if loss_type == "softmax":
+                current_sample = fluid.layers.clip(
+                    current_sample, min=-1.0, max=0.99999)
+                # quantized have values in [0, num_channels)
+                quantized = fluid.layers.cast(
+                    (current_sample + 1.0) / 2.0 * self.config.num_channels,
+                    dtype="int64")
+                audio_input = self.embedding_fc(quantized)
+            elif loss_type == "mix-gaussian-pdf":
+                audio_input = self.embedding_fc(current_sample)
+            else:
+                raise ValueError(
+                    "loss_type {} is unsupported!".format(loss_type))
+
+            # [bs, 128, 1, 1]
+            audio_input = fluid.layers.unsqueeze(fluid.layers.transpose(audio_input, perm=[0, 2, 1]), 2)
+            # [bs, 80]
+            cond_input = conditioner[:, i, :]
+            # [bs, 80, 1, 1]
+            cond_input = fluid.layers.reshape(
+                cond_input, cond_input.shape + [1, 1])
+
+            skip = None
+            for layer in self.dilated_causal_convs:
+                audio_input, skip = layer.add_input(audio_input, skip, cond_input)
+            
+            # [bs, 1, 128]
+            skip = fluid.layers.transpose(fluid.layers.squeeze(skip, [2]), perm=[0, 2, 1])
+            # [bs, 1, 3]
+            mix_parameters = self.fc3(self.fc2(self.fc1(skip)))
+            if loss_type == "softmax":
+                sample = self.sample_softmax(mix_parameters)
+            elif loss_type == "mix-gaussian-pdf":
+                sample = self.sample_mix_gaussian(mix_parameters)
+            else:
+                raise ValueError(
+                    "loss_type {} is unsupported!".format(loss_type))
+            audio_samples.append(sample)
+            # [bs]
+            current_sample = audio_samples[-1]
+            # [bs, 1, 1]
+            current_sample = fluid.layers.reshape(current_sample,
+                current_sample.shape + [1, 1])
+
+        # syn_audio: (num_samples,)
+        syn_audio = fluid.layers.concat(audio_samples, axis=0).numpy()
+
+        return syn_audio        
+
+    def start_new_sequence(self):
+        for layer in self.sublayers():
+            if isinstance(layer, weight_norm.Conv1D):
+                layer.start_new_sequence()
+
+    def save(self, iteration):
+        utils.save_latest_parameters(self.checkpoint_dir, iteration,
+                                     self.wavenet, self.optimizer)
+        utils.save_latest_checkpoint(self.checkpoint_dir, iteration)