Merge branch 'master' into 'master'

Adding WaveNet model verified on LJSpeech dataset

See merge request !3

.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/data/sampler.py

@@ -163,6 +163,35 @@ class WeightedRandomSampler(Sampler):
         return self.num_samples
 
 
+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(np.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 BatchSampler(Sampler):
     r"""Wraps another sampler to yield a mini-batch of indices.
     Args:
@@ -206,4 +235,4 @@ class BatchSampler(Sampler):
         if self.drop_last:
             return len(self.sampler) // self.batch_size
         else:
-            return (len(self.sampler) + self.batch_size - 1) // self.batch_size
\ No newline at end of file
+            return (len(self.sampler) + self.batch_size - 1) // self.batch_size

parakeet/models/wavenet/README.md

+# WaveNet with Paddle Fluid
+
+Paddle fluid implementation of WaveNet, a deep generative model of raw audio waveforms.
+WaveNet model is originally proposed in [WaveNet: A Generative Model for Raw Audio](https://arxiv.org/abs/1609.03499).
+Our implementation is based on the WaveNet architecture described in [ClariNet: Parallel Wave Generation in End-to-End Text-to-Speech](https://arxiv.org/abs/1807.07281) and can provide various output distributions, including single Gaussian, mixture of Gaussian, and softmax with linearly quantized channels. 
+
+We implement WaveNet model in paddle fluid with dynamic graph, which is convenient for flexible network architectures.
+
+## Project Structure
+```text
+├── configs                 # yaml configuration files of preset model hyperparameters
+├── data.py                 # dataset and dataloader settings for LJSpeech
+├── slurm.py                # optional slurm helper functions if you use slurm to train model
+├── synthesis.py            # script for speech synthesis
+├── train.py                # script for model training
+├── utils.py                # helper functions for e.g., model checkpointing
+├── wavenet.py              # WaveNet model high level APIs
+└── wavenet_modules.py      # WaveNet model implementation
+```
+
+## Usage
+
+There are many hyperparameters to be tuned depending on the specification of model and dataset you are working on. Hyperparameters that are known to work good for the LJSpeech dataset are provided as yaml files in `./configs/` folder. Specifically, we provide `wavenet_ljspeech_single_gaussian.yaml`, `wavenet_ljspeech_mix_gaussian.yaml`, and `wavenet_ljspeech_softmax.yaml` config files for WaveNet with single Gaussian, 10-component mixture of Gaussians, and softmax (with 2048 linearly quantized channels) output distributions, respectively.
+
+Note that `train.py` and `synthesis.py` all accept a `--config` parameter. To ensure consistency, you should use the same config yaml file for both training and synthesizing. You can also overwrite these preset hyperparameters with command line by updating parameters after `--config`. For example `--config=${yaml} --batch_size=8 --layers=20` can overwrite the corresponding hyperparameters in the `${yaml}` config file. For more details about these hyperparameters, check `utils.add_config_options_to_parser`.
+
+Note that you also need to specify some additional parameters for `train.py` and `synthesis.py`, and the details can be found in `train.add_options_to_parser` and `synthesis.add_options_to_parser`, respectively.
+
+### Dataset
+
+Download and unzip [LJSpeech](https://keithito.com/LJ-Speech-Dataset/).
+
+```bash
+wget https://data.keithito.com/data/speech/LJSpeech-1.1.tar.bz2
+tar xjvf LJSpeech-1.1.tar.bz2
+```
+
+In this example, assume that the path of unzipped LJSpeech dataset is `./data/LJSpeech-1.1`.
+
+### Train on single GPU
+
+```bash
+export PYTHONPATH="${PYTHONPATH}:${PWD}/../../.."
+export CUDA_VISIBLE_DEVICES=0
+python -u train.py --config=${yaml} \
+    --root=./data/LJSpeech-1.1 \
+    --name=${ModelName} --batch_size=4 \
+    --parallel=false --use_gpu=true
+```
+
+#### Save and Load checkpoints
+
+Our model will save model parameters as checkpoints in `./runs/wavenet/${ModelName}/checkpoint/` every 10000 iterations by default.
+The saved checkpoint will have the format of `step-${iteration_number}.pdparams` for model parameters and `step-${iteration_number}.pdopt` for optimizer parameters. 
+
+There are three ways to load a checkpoint and resume training (take an example that you want to load a 500000-iteration checkpoint):
+1. Use `--checkpoint=./runs/wavenet/${ModelName}/checkpoint/step-500000` to provide a specific path to load. Note that you only need to provide the base name of the parameter file, which is `step-500000`, no extension name `.pdparams` or `.pdopt` is needed. 
+2. Use `--iteration=500000`.
+3. If you don't specify either `--checkpoint` or `--iteration`, the model will automatically load the latest checkpoint in `./runs/wavenet/${ModelName}/checkpoint`.
+
+### Train on multiple GPUs
+
+```bash
+export PYTHONPATH="${PYTHONPATH}:${PWD}/../../.."
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+python -u -m paddle.distributed.launch train.py \
+    --config=${yaml} \
+    --root=./data/LJSpeech-1.1 \
+    --name=${ModelName} --parallel=true --use_gpu=true
+```
+
+Use `export CUDA_VISIBLE_DEVICES=0,1,2,3` to set the GPUs that you want to use to be visible. Then the `paddle.distributed.launch` module will use these visible GPUs to do data parallel training in multiprocessing mode.
+
+### Monitor with Tensorboard
+
+By default, the logs are saved in `./runs/wavenet/${ModelName}/logs/`. You can monitor logs by tensorboard.
+
+```bash
+tensorboard --logdir=${log_dir} --port=8888
+```
+
+### Synthesize from a checkpoint
+
+Check the [Save and load checkpoint](#save-and-load-checkpoints) section on how to load a specific checkpoint.
+The following example will automatically load the latest checkpoint:
+
+```bash
+export PYTHONPATH="${PYTHONPATH}:${PWD}/../../.."
+export CUDA_VISIBLE_DEVICES=0
+python -u synthesis.py --config=${yaml} \
+    --root=./data/LJSpeech-1.1 \
+    --name=${ModelName} --use_gpu=true \
+    --output=./syn_audios \
+    --sample=${SAMPLE} 
+```
+
+In this example, `--output` specifies where to save the synthesized audios and `--sample` specifies which sample in the valid dataset (a split from the whole LJSpeech dataset, by default contains the first 16 audio samples) to synthesize based on the mel-spectrograms computed from the ground truth sample audio, e.g., `--sample=0` means to synthesize the first audio in the valid dataset.

parakeet/models/wavenet/configs/wavenet_ljspeech_mix_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: 10
+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/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/configs/wavenet_ljspeech_softmax.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: softmax
+num_channels: 2048
+
+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 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 DistributedSampler, BatchSampler
+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(np.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 = int(np.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 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/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
+import paddle.fluid.dygraph as dg
+from paddle import fluid
+
+import utils
+from data import LJSpeech
+from wavenet_modules import WaveNetModule
+
+
+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
+
+        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 = []
+        sample_audios = []
+        start_time = time.time()
+        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)
+
+    @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)
+
+    def save(self, iteration):
+        utils.save_latest_parameters(self.checkpoint_dir, iteration,
+                                     self.wavenet, self.optimizer)
+        utils.save_latest_checkpoint(self.checkpoint_dir, iteration)

parakeet/models/wavenet/wavenet_modules.py

+import itertools
+
+import numpy as np
+import paddle.fluid.dygraph as dg
+from paddle import fluid
+from parakeet.modules import conv, modules
+
+
+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 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(
+                modules.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 = modules.Embedding(
+                self.full_name(),
+                num_embeddings=config.num_channels,
+                embed_dim=config.residual_channels,
+                std=0.1)
+        elif config.loss_type == "mix-gaussian-pdf":
+            self.embedding_fc = modules.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(
+                modules.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 = modules.FC(
+            self.full_name(),
+            in_features=config.residual_channels,
+            size=config.skip_channels,
+            num_flatten_dims=2,
+            relu=True,
+            act="relu")
+
+        self.fc2 = modules.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 = modules.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 = modules.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, len, 3 * num_mixtures]
+        # to [bs * len, 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 = 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, len, 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)
+
+        return loss
+
+    def mixture_density_loss(self, targets, mix_parameters, log_scale_min):
+        # targets: [bs, len]
+        # mix_params: [bs, len, 3 * num_mixture]
+        targets = targets[:, self.context_size:]
+        mix_parameters = mix_parameters[:, self.context_size:, :]
+
+        # log_s: [bs, len, 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 = 1.0 / np.sqrt(2.0 * np.pi) * inv_s * exponent
+        pdf_x = pi * pdf_x
+        # pdf_x: [bs, len]
+        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):
+        # 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, len]
+        input_audios = audios[:, :-1]
+        target_audios = audios[:, 1:]
+        # conditioner: [bs, len, mel_bands]
+        conditioner = conditioner[:, 1:, :]
+
+        loss_type = self.config.loss_type
+
+        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, len]
+        layer_input = fluid.layers.unsqueeze(
+            fluid.layers.transpose(layer_input, perm=[0, 2, 1]), 2)
+        # conditioner: [bs, mel_bands, 1, len]
+        conditioner = fluid.layers.unsqueeze(
+            fluid.layers.transpose(conditioner, perm=[0, 2, 1]), 2)
+
+        skip = None
+        for i, layer in enumerate(self.dilated_causal_convs):
+            # layer_input: [bs, res_channel, 1, len]
+            # skip: [bs, res_channel, 1, len]
+            layer_input, skip = layer(layer_input, skip, conditioner)
+
+        # Reshape skip to [bs, len, res_channel]
+        skip = fluid.layers.transpose(
+            fluid.layers.squeeze(skip, [2]), perm=[0, 2, 1])
+        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))
+
+        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))
+
+        return loss, sample_audios
+
+    def synthesize(self, mels):
+        self.start_new_sequence()
+        bs, n_frames, mel_bands = mels.shape 
+        conditioner = self.conditioner(mels)
+        time_steps = conditioner.shape[1]
+
+        print("input mels shape", mels.shape)
+        print("Total synthesis 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.
+            # audio_input: [bs, 1, channel]
+            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, channel, 1, 1]
+            audio_input = fluid.layers.unsqueeze(
+                fluid.layers.transpose(audio_input, perm=[0, 2, 1]), 2)
+            # [bs, mel_bands]
+            cond_input = conditioner[:, i, :]
+            # [bs, mel_bands, 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, channel]
+            skip = fluid.layers.transpose(
+                fluid.layers.squeeze(skip, [2]), perm=[0, 2, 1])
+            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, conv.Conv1D):
+                layer.start_new_sequence()

parakeet/modules/modules.py

@@ -26,6 +26,7 @@ def FC(name_scope,
        in_features,
        size,
        num_flatten_dims=1,
+       relu=False,
        dropout=0.0,
        epsilon=1e-30,
        act=None,
@@ -39,7 +40,11 @@ def FC(name_scope,
     # stds
     if isinstance(in_features, int):
         in_features = [in_features]
+
     stds = [np.sqrt((1 - 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
     ]
@@ -456,3 +461,152 @@ class PositionEmbedding(dg.Layer):
             return out
         else:
             raise Exception("Then you can just use position rate at init")
+
+
+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 Conv1D_GU
+                layer, where B means batch_size, C_in means the input channels
+                T means input time steps.
+            skip (Variable): Shape(B, C_in, 1, T), skip connection.
+            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 Conv1D_GU 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 Conv1D_GU, where
+                C_out means the output channels of Conv1D_GU.
+            skip (Variable): Shape(B, C_out, 1, T), skip connection.
+        """
+        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)
+            skip: shape(B, num_filters, 1, time_steps), skip connection
+            conditioner: shape(B, conditioner_dim, 1, time_steps)
+        Outputs:
+            x: shape(B, num_filters, 1, time_steps), where time_steps = 1
+            skip: skip connection, same shape as x
+        """
+        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