1. using SentencePiece, Pretrained BPEmb

2. Using 40 Mel channels with 4000Hz upperbound

.gitignore

@@ -0,0 +1,147 @@
+.env
+env/
+jupyter.json
+run*/
+filelists/
+# Created by https://www.gitignore.io/api/macos
+# Edit at https://www.gitignore.io/?templates=macos
+
+### macOS ###
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+# End of https://www.gitignore.io/api/macos
+
+# Created by https://www.gitignore.io/api/python
+# Edit at https://www.gitignore.io/?templates=python
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# Mr Developer
+.mr.developer.cfg
+.project
+.pydevproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# End of https://www.gitignore.io/api/python

.gitmodules

@@ -0,0 +1,4 @@
+[submodule "waveglow"]
+	path = waveglow
+	url = https://github.com/NVIDIA/waveglow
+	branch = master

Dockerfile

@@ -1,2 +1,10 @@
-FROM pytorch/pytorch:0.4_cuda9_cudnn7
+FROM pytorch/pytorch:nightly-devel-cuda10.0-cudnn7
+ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:${PATH}
+
+RUN apt-get update -y
+
 RUN pip install numpy scipy matplotlib librosa==0.6.0 tensorflow tensorboardX inflect==0.2.5 Unidecode==1.0.22 jupyter
+
+ADD apex /apex/
+WORKDIR /apex/
+RUN pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .

Makefile

@@ -0,0 +1,6 @@
+.PHONY: clean clean-test clean-pyc clean-build docs help common.mk
+.DEFAULT_GOAL := help
+
+
+notebook:
+	jupyter lab --ip=0.0.0.0 --no-browser --NotebookApp.token='${JUPYTER_TOKEN}'

README.md

@@ -3,15 +3,16 @@
 PyTorch implementation of [Natural TTS Synthesis By Conditioning
 Wavenet On Mel Spectrogram Predictions](https://arxiv.org/pdf/1712.05884.pdf). 
 
-This implementation includes **distributed** and **fp16** support
+This implementation includes **distributed** and **automatic mixed precision** support
 and uses the [LJSpeech dataset](https://keithito.com/LJ-Speech-Dataset/).
 
-Distributed and FP16 support relies on work by Christian Sarofeen and NVIDIA's
-[Apex Library](https://github.com/nvidia/apex).
+Distributed and Automatic Mixed Precision support relies on NVIDIA's [Apex] and [AMP].
+
+Visit our [website] for audio samples using our published [Tacotron 2] and
+[WaveGlow] models.
 
 ![Alignment, Predicted Mel Spectrogram, Target Mel Spectrogram](tensorboard.png)
 
-Visit our [website] for audio samples.
 
 ## Pre-requisites
 1. NVIDIA GPU + CUDA cuDNN
@@ -20,17 +21,26 @@ Visit our [website] for audio samples.
 1. Download and extract the [LJ Speech dataset](https://keithito.com/LJ-Speech-Dataset/)
 2. Clone this repo: `git clone https://github.com/NVIDIA/tacotron2.git`
 3. CD into this repo: `cd tacotron2`
-4. Update .wav paths: `sed -i -- 's,DUMMY,ljs_dataset_folder/wavs,g' filelists/*.txt`
+4. Initialize submodule: `git submodule init; git submodule update`
+5. Update .wav paths: `sed -i -- 's,DUMMY,ljs_dataset_folder/wavs,g' filelists/*.txt`
     - Alternatively, set `load_mel_from_disk=True` in `hparams.py` and update mel-spectrogram paths 
-5. Install [PyTorch 1.0]
-6. Install python requirements or build docker image 
+6. Install [PyTorch 1.0]
+7. Install [Apex]
+8. Install python requirements or build docker image 
     - Install python requirements: `pip install -r requirements.txt`
 
 ## Training
 1. `python train.py --output_directory=outdir --log_directory=logdir`
 2. (OPTIONAL) `tensorboard --logdir=outdir/logdir`
 
-## Multi-GPU (distributed) and FP16 Training
+## Training using a pre-trained model
+Training using a pre-trained model can lead to faster convergence  
+By default, the dataset dependent text embedding layers are [ignored]
+
+1. Download our published [Tacotron 2] model
+2. `python train.py --output_directory=outdir --log_directory=logdir -c tacotron2_statedict.pt --warm_start`
+
+## Multi-GPU (distributed) and Automatic Mixed Precision Training
 1. `python -m multiproc train.py --output_directory=outdir --log_directory=logdir --hparams=distributed_run=True,fp16_run=True`
 
 ## Inference demo
@@ -62,7 +72,10 @@ We are thankful to the Tacotron 2 paper authors, specially Jonathan Shen, Yuxuan
 Wang and Zongheng Yang.
 
 
-[WaveGlow]: https://drive.google.com/file/d/1cjKPHbtAMh_4HTHmuIGNkbOkPBD9qwhj/view?usp=sharing
+[WaveGlow]: https://drive.google.com/file/d/1WsibBTsuRg_SF2Z6L6NFRTT-NjEy1oTx/view?usp=sharing
 [Tacotron 2]: https://drive.google.com/file/d/1c5ZTuT7J08wLUoVZ2KkUs_VdZuJ86ZqA/view?usp=sharing
 [pytorch 1.0]: https://github.com/pytorch/pytorch#installation
 [website]: https://nv-adlr.github.io/WaveGlow
+[ignored]: https://github.com/NVIDIA/tacotron2/blob/master/hparams.py#L22
+[Apex]: https://github.com/nvidia/apex
+[AMP]: https://github.com/NVIDIA/apex/tree/master/apex/amp

common.mk

@@ -0,0 +1,27 @@
+define BROWSER_PYSCRIPT
+import os, webbrowser, sys
+
+try:
+	from urllib import pathname2url
+except:
+	from urllib.request import pathname2url
+
+webbrowser.open("file://" + pathname2url(os.path.abspath(sys.argv[1])))
+endef
+export BROWSER_PYSCRIPT
+
+define PRINT_HELP_PYSCRIPT
+import re, sys
+
+for line in sys.stdin:
+	match = re.match(r'^([a-zA-Z_-]+):.*?## (.*)$$', line)
+	if match:
+		target, help = match.groups()
+		print("%-20s %s" % (target, help))
+endef
+export PRINT_HELP_PYSCRIPT
+
+BROWSER := python -c "$$BROWSER_PYSCRIPT"
+
+help:  ##  make TARGET forwards the TARGET to sub packages
+	@cat $(MAKEFILE_LIST) | python -c "$$PRINT_HELP_PYSCRIPT"

data_utils.py

@@ -5,8 +5,8 @@ import torch.utils.data
 
 import layers
 from utils import load_wav_to_torch, load_filepaths_and_text
-from text import text_to_sequence
-
+# from text import text_to_sequence
+from text_codec import text_to_sequence
 
 class TextMelLoader(torch.utils.data.Dataset):
     """

distributed.py

@@ -1,6 +1,7 @@
 import torch
 import torch.distributed as dist
 from torch.nn.modules import Module
+from torch.autograd import Variable
 
 def _flatten_dense_tensors(tensors):
     """Flatten dense tensors into a contiguous 1D buffer. Assume tensors are of
@@ -139,7 +140,7 @@ def apply_gradient_allreduce(module):
                 buckets = {}
                 for param in module.parameters():
                     if param.requires_grad and param.grad is not None:
-                        tp = type(param.data)
+                        tp = param.data.dtype
                         if tp not in buckets:
                             buckets[tp] = []
                         buckets[tp].append(param)
@@ -161,12 +162,12 @@ def apply_gradient_allreduce(module):
 
         for param in list(module.parameters()):
             def allreduce_hook(*unused):
-                param._execution_engine.queue_callback(allreduce_params)
+                Variable._execution_engine.queue_callback(allreduce_params)
             if param.requires_grad:
                 param.register_hook(allreduce_hook)
-   
+
         def set_needs_reduction(self, input, output):
             self.needs_reduction = True
-        
+
         module.register_forward_hook(set_needs_reduction)
         return module

fp16_optimizer.py

@@ -1,381 +0,0 @@
-import torch
-from torch import nn
-from torch.autograd import Variable
-from torch.nn.parameter import Parameter
-from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
-
-from loss_scaler import DynamicLossScaler, LossScaler
-
-FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor)
-HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor)
-
-def conversion_helper(val, conversion):
-    """Apply conversion to val. Recursively apply conversion if `val` is a nested tuple/list structure."""
-    if not isinstance(val, (tuple, list)):
-        return conversion(val)
-    rtn =  [conversion_helper(v, conversion) for v in val]
-    if isinstance(val, tuple):
-        rtn = tuple(rtn)
-    return rtn
-
-def fp32_to_fp16(val):
-    """Convert fp32 `val` to fp16"""
-    def half_conversion(val):
-        val_typecheck = val
-        if isinstance(val_typecheck, (Parameter, Variable)):
-            val_typecheck = val.data
-        if isinstance(val_typecheck, FLOAT_TYPES):
-            val = val.half()
-        return val
-    return conversion_helper(val, half_conversion)
-
-def fp16_to_fp32(val):
-    """Convert fp16 `val` to fp32"""
-    def float_conversion(val):
-        val_typecheck = val
-        if isinstance(val_typecheck, (Parameter, Variable)):
-            val_typecheck = val.data
-        if isinstance(val_typecheck, HALF_TYPES):
-            val = val.float()
-        return val
-    return conversion_helper(val, float_conversion)
-
-class FP16_Module(nn.Module):
-    def __init__(self, module):
-        super(FP16_Module, self).__init__()
-        self.add_module('module', module.half())
-
-    def forward(self, *inputs, **kwargs):
-        return fp16_to_fp32(self.module(*(fp32_to_fp16(inputs)), **kwargs))
-
-class FP16_Optimizer(object):
-    """
-    FP16_Optimizer is designed to wrap an existing PyTorch optimizer,
-    and enable an fp16 model to be trained using a master copy of fp32 weights.
-
-    Args:
-        optimizer (torch.optim.optimizer):  Existing optimizer containing initialized fp16 parameters.  Internally, FP16_Optimizer replaces the passed optimizer's fp16 parameters with new fp32 parameters copied from the original ones.  FP16_Optimizer also stores references to the original fp16 parameters, and updates these fp16 parameters from the master fp32 copy after each step.
-        static_loss_scale (float, optional, default=1.0):  Loss scale used internally to scale fp16 gradients computed by the model.  Scaled gradients will be copied to fp32, then downscaled before being applied to the fp32 master params, so static_loss_scale should not affect learning rate.
-        dynamic_loss_scale (bool, optional, default=False):  Use dynamic loss scaling.  If True, this will override any static_loss_scale option.
-
-    """
-
-    def __init__(self, optimizer, static_loss_scale=1.0, dynamic_loss_scale=False):
-        if not torch.cuda.is_available:
-            raise SystemError('Cannot use fp16 without CUDA')
-
-        self.fp16_param_groups = []
-        self.fp32_param_groups = []
-        self.fp32_flattened_groups = []
-        for i, param_group in enumerate(optimizer.param_groups):
-            print("FP16_Optimizer processing param group {}:".format(i))
-            fp16_params_this_group = []
-            fp32_params_this_group = []
-            for param in param_group['params']:
-                if param.requires_grad:
-                    if param.type() == 'torch.cuda.HalfTensor':
-                        print("FP16_Optimizer received torch.cuda.HalfTensor with {}"
-                              .format(param.size()))
-                        fp16_params_this_group.append(param)
-                    elif param.type() == 'torch.cuda.FloatTensor':
-                        print("FP16_Optimizer received torch.cuda.FloatTensor with {}"
-                              .format(param.size()))
-                        fp32_params_this_group.append(param)
-                    else:
-                        raise TypeError("Wrapped parameters must be either "
-                                        "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
-                                        "Received {}".format(param.type()))
-
-            fp32_flattened_this_group = None
-            if len(fp16_params_this_group) > 0:
-                fp32_flattened_this_group = _flatten_dense_tensors(
-                    [param.detach().data.clone().float() for param in fp16_params_this_group])
-
-                fp32_flattened_this_group = Variable(fp32_flattened_this_group, requires_grad = True)
-
-                fp32_flattened_this_group.grad = fp32_flattened_this_group.new(
-                    *fp32_flattened_this_group.size())
-
-            # python's lovely list concatenation via +
-            if fp32_flattened_this_group is not None:
-                param_group['params'] = [fp32_flattened_this_group] + fp32_params_this_group
-            else:
-                param_group['params'] = fp32_params_this_group
-
-            self.fp16_param_groups.append(fp16_params_this_group)
-            self.fp32_param_groups.append(fp32_params_this_group)
-            self.fp32_flattened_groups.append(fp32_flattened_this_group)
-
-        # print("self.fp32_flattened_groups = ", self.fp32_flattened_groups)
-        # print("self.fp16_param_groups = ", self.fp16_param_groups)
-
-        self.optimizer = optimizer.__class__(optimizer.param_groups)
-
-        # self.optimizer.load_state_dict(optimizer.state_dict())
-
-        self.param_groups = self.optimizer.param_groups
-
-        if dynamic_loss_scale:
-            self.dynamic_loss_scale = True
-            self.loss_scaler = DynamicLossScaler()
-        else:
-            self.dynamic_loss_scale = False
-            self.loss_scaler = LossScaler(static_loss_scale)
-
-        self.overflow = False
-        self.first_closure_call_this_step = True
-
-    def zero_grad(self):
-        """
-        Zero fp32 and fp16 parameter grads.
-        """
-        self.optimizer.zero_grad()
-        for fp16_group in self.fp16_param_groups:
-            for param in fp16_group:
-                if param.grad is not None:
-                    param.grad.detach_() # This does appear in torch.optim.optimizer.zero_grad(),
-                                         # but I'm not sure why it's needed.
-                    param.grad.zero_()
-
-    def _check_overflow(self):
-        params = []
-        for group in self.fp16_param_groups:
-            for param in group:
-                params.append(param)
-        for group in self.fp32_param_groups:
-            for param in group:
-                params.append(param)
-        self.overflow = self.loss_scaler.has_overflow(params)
-
-    def _update_scale(self, has_overflow=False):
-        self.loss_scaler.update_scale(has_overflow)
-
-    def _copy_grads_fp16_to_fp32(self):
-        for fp32_group, fp16_group in zip(self.fp32_flattened_groups, self.fp16_param_groups):
-            if len(fp16_group) > 0:
-                # This might incur one more deep copy than is necessary.
-                fp32_group.grad.data.copy_(
-                    _flatten_dense_tensors([fp16_param.grad.data for fp16_param in fp16_group]))
-
-    def _downscale_fp32(self):
-        if self.loss_scale != 1.0:
-            for param_group in self.optimizer.param_groups:
-                for param in param_group['params']:
-                    param.grad.data.mul_(1./self.loss_scale)
-
-    def clip_fp32_grads(self, clip=-1):
-        if not self.overflow:
-            fp32_params = []
-            for param_group in self.optimizer.param_groups:
-                for param in param_group['params']:
-                    fp32_params.append(param)
-            if clip > 0:
-                return torch.nn.utils.clip_grad_norm(fp32_params, clip)
-
-    def _copy_params_fp32_to_fp16(self):
-        for fp16_group, fp32_group in zip(self.fp16_param_groups, self.fp32_flattened_groups):
-            if len(fp16_group) > 0:
-                for fp16_param, fp32_data in zip(fp16_group,
-                    _unflatten_dense_tensors(fp32_group.data, fp16_group)):
-                    fp16_param.data.copy_(fp32_data)
-
-    def state_dict(self):
-        """
-        Returns a dict containing the current state of this FP16_Optimizer instance.
-        This dict contains attributes of FP16_Optimizer, as well as the state_dict
-        of the contained Pytorch optimizer.
-
-        Untested.
-        """
-        state_dict = {}
-        state_dict['loss_scaler'] = self.loss_scaler
-        state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale
-        state_dict['overflow'] = self.overflow
-        state_dict['first_closure_call_this_step'] = self.first_closure_call_this_step
-        state_dict['optimizer_state_dict'] = self.optimizer.state_dict()
-        return state_dict
-
-    def load_state_dict(self, state_dict):
-        """
-        Loads a state_dict created by an earlier call to state_dict.
-
-        Untested.
-        """
-        self.loss_scaler = state_dict['loss_scaler']
-        self.dynamic_loss_scale = state_dict['dynamic_loss_scale']
-        self.overflow = state_dict['overflow']
-        self.first_closure_call_this_step = state_dict['first_closure_call_this_step']
-        self.optimizer.load_state_dict(state_dict['optimizer_state_dict'])
-
-    def step(self, closure=None): # could add clip option.
-        """
-        If no closure is supplied, step should be called after fp16_optimizer_obj.backward(loss).
-        step updates the fp32 master copy of parameters using the optimizer supplied to
-        FP16_Optimizer's constructor, then copies the updated fp32 params into the fp16 params
-        originally referenced by Fp16_Optimizer's constructor, so the user may immediately run
-        another forward pass using their model.
-
-        If a closure is supplied, step may be called without a prior call to self.backward(loss).
-        However, the user should take care that any loss.backward() call within the closure
-        has been replaced by fp16_optimizer_obj.backward(loss).
-
-        Args:
-           closure (optional):  Closure that will be supplied to the underlying optimizer originally passed to FP16_Optimizer's constructor.  closure should call zero_grad on the FP16_Optimizer object, compute the loss, call .backward(loss), and return the loss.
-
-        Closure example::
-
-            # optimizer is assumed to be an FP16_Optimizer object, previously constructed from an
-            # existing pytorch optimizer.
-            for input, target in dataset:
-                def closure():
-                    optimizer.zero_grad()
-                    output = model(input)
-                    loss = loss_fn(output, target)
-                    optimizer.backward(loss)
-                    return loss
-                optimizer.step(closure)
-
-        .. note::
-            The only changes that need to be made compared to
-            `ordinary optimizer closures`_ are that "optimizer" itself should be an instance of
-            FP16_Optimizer, and that the call to loss.backward should be replaced by
-            optimizer.backward(loss).
-
-        .. warning::
-            Currently, calling step with a closure is not compatible with dynamic loss scaling.
-
-        .. _`ordinary optimizer closures`:
-            http://pytorch.org/docs/master/optim.html#optimizer-step-closure
-        """
-        if closure is not None and isinstance(self.loss_scaler, DynamicLossScaler):
-            raise TypeError("Using step with a closure is currently not "
-                            "compatible with dynamic loss scaling.")
-
-        scale = self.loss_scaler.loss_scale
-        self._update_scale(self.overflow)
-
-        if self.overflow:
-            print("OVERFLOW! Skipping step. Attempted loss scale: {}".format(scale))
-            return
-
-        if closure is not None:
-            self._step_with_closure(closure)
-        else:
-            self.optimizer.step()
-
-        self._copy_params_fp32_to_fp16()
-
-        return
-
-    def _step_with_closure(self, closure):
-        def wrapped_closure():
-            if self.first_closure_call_this_step:
-                """
-                We expect that the fp16 params are initially fresh on entering self.step(),
-                so _copy_params_fp32_to_fp16() is unnecessary the first time wrapped_closure()
-                is called within self.optimizer.step().
-                """
-                self.first_closure_call_this_step = False
-            else:
-                """
-                If self.optimizer.step() internally calls wrapped_closure more than once,
-                it may update the fp32 params after each call.  However, self.optimizer
-                doesn't know about the fp16 params at all.  If the fp32 params get updated,
-                we can't rely on self.optimizer to refresh the fp16 params.  We need
-                to handle that manually:
-                """
-                self._copy_params_fp32_to_fp16()
-
-            """
-            Our API expects the user to give us ownership of the backward() call by
-            replacing all calls to loss.backward() with optimizer.backward(loss).
-            This requirement holds whether or not the call to backward() is made within
-            a closure.
-            If the user is properly calling optimizer.backward(loss) within "closure,"
-            calling closure() here will give the fp32 master params fresh gradients
-            for the optimizer to play with,
-            so all wrapped_closure needs to do is call closure() and return the loss.
-            """
-            temp_loss = closure()
-            return temp_loss
-
-        self.optimizer.step(wrapped_closure)
-
-        self.first_closure_call_this_step = True
-
-    def backward(self, loss, update_fp32_grads=True):
-        """
-        fp16_optimizer_obj.backward performs the following conceptual operations:
-
-        fp32_loss = loss.float() (see first Note below)
-
-        scaled_loss = fp32_loss*loss_scale
-
-        scaled_loss.backward(), which accumulates scaled gradients into the .grad attributes of the
-        fp16 model's leaves.
-
-        fp16 grads are then copied to the stored fp32 params' .grad attributes (see second Note).
-
-        Finally, fp32 grads are divided by loss_scale.
-
-        In this way, after fp16_optimizer_obj.backward, the fp32 parameters have fresh gradients,
-        and fp16_optimizer_obj.step may be called.
-
-        .. note::
-            Converting the loss to fp32 before applying the loss scale provides some
-            additional safety against overflow if the user has supplied an fp16 value.
-            However, for maximum overflow safety, the user should
-            compute the loss criterion (MSE, cross entropy, etc) in fp32 before supplying it to
-            fp16_optimizer_obj.backward.
-
-        .. note::
-            The gradients found in an fp16 model's leaves after a call to
-            fp16_optimizer_obj.backward should not be regarded as valid in general,
-            because it's possible
-            they have been scaled (and in the case of dynamic loss scaling,
-            the scale factor may silently change over time).
-            If the user wants to inspect gradients after a call to fp16_optimizer_obj.backward,
-            he/she should query the .grad attribute of FP16_Optimizer's stored fp32 parameters.
-
-        Args:
-            loss:  The loss output by the user's model.  loss may be either float or half (but see first Note above).
-            update_fp32_grads (bool, optional, default=True):  Option to copy fp16 grads to fp32 grads on this call.  By setting this to False, the user can delay this copy, which is useful to eliminate redundant fp16->fp32 grad copies if fp16_optimizer_obj.backward is being called on multiple losses in one iteration.  If set to False, the user becomes responsible for calling fp16_optimizer_obj.update_fp32_grads before calling fp16_optimizer_obj.step.
-
-        Example::
-
-            # Ordinary operation:
-            optimizer.backward(loss)
-
-            # Naive operation with multiple losses (technically valid, but less efficient):
-            # fp32 grads will be correct after the second call,  but
-            # the first call incurs an unnecessary fp16->fp32 grad copy.
-            optimizer.backward(loss1)
-            optimizer.backward(loss2)
-
-            # More efficient way to handle multiple losses:
-            # The fp16->fp32 grad copy is delayed until fp16 grads from all
-            # losses have been accumulated.
-            optimizer.backward(loss1, update_fp32_grads=False)
-            optimizer.backward(loss2, update_fp32_grads=False)
-            optimizer.update_fp32_grads()
-        """
-        self.loss_scaler.backward(loss.float())
-        if update_fp32_grads:
-            self.update_fp32_grads()
-
-    def update_fp32_grads(self):
-        """
-        Copy the .grad attribute from stored references to fp16 parameters to
-        the .grad attribute of the master fp32 parameters that are directly
-        updated by the optimizer.  :attr:`update_fp32_grads` only needs to be called if
-        fp16_optimizer_obj.backward was called with update_fp32_grads=False.
-        """
-        if self.dynamic_loss_scale:
-            self._check_overflow()
-            if self.overflow: return
-        self._copy_grads_fp16_to_fp32()
-        self._downscale_fp32()
-
-    @property
-    def loss_scale(self):
-        return self.loss_scaler.loss_scale

hparams.py

@@ -16,16 +16,17 @@ def create_hparams(hparams_string=None, verbose=False):
         fp16_run=False,
         distributed_run=False,
         dist_backend="nccl",
-        dist_url="file://distributed.dpt",
+        dist_url="tcp://localhost:54321",
         cudnn_enabled=True,
         cudnn_benchmark=False,
+        ignore_layers=['embedding.weight'],
 
         ################################
         # Data Parameters             #
         ################################
         load_mel_from_disk=False,
-        training_files='filelists/ljs_audio22khz_text_train_filelist.txt',
-        validation_files='filelists/ljs_audio22khz_text_val_filelist.txt',
+        training_files='filelists/ljs_audio_text_train_filelist.txt',
+        validation_files='filelists/ljs_audio_text_val_filelist.txt',
         text_cleaners=['english_cleaners'],
 
         ################################
@@ -36,14 +37,14 @@ def create_hparams(hparams_string=None, verbose=False):
         filter_length=1024,
         hop_length=256,
         win_length=1024,
-        n_mel_channels=80,
+        n_mel_channels=40,
         mel_fmin=0.0,
-        mel_fmax=8000.0,
+        mel_fmax=4000.0,
 
         ################################
         # Model Parameters             #
         ################################
-        n_symbols=len(symbols),
+        n_symbols=1000,#len(symbols),
         symbols_embedding_dim=512,
 
         # Encoder parameters

logger.py

@@ -1,5 +1,5 @@
 import random
-import torch.nn.functional as F
+import torch
 from tensorboardX import SummaryWriter
 from plotting_utils import plot_alignment_to_numpy, plot_spectrogram_to_numpy
 from plotting_utils import plot_gate_outputs_to_numpy
@@ -44,5 +44,5 @@ class Tacotron2Logger(SummaryWriter):
             "gate",
             plot_gate_outputs_to_numpy(
                 gate_targets[idx].data.cpu().numpy(),
-                F.sigmoid(gate_outputs[idx]).data.cpu().numpy()),
+                torch.sigmoid(gate_outputs[idx]).data.cpu().numpy()),
             iteration)

model.py

@@ -5,7 +5,6 @@ from torch import nn
 from torch.nn import functional as F
 from layers import ConvNorm, LinearNorm
 from utils import to_gpu, get_mask_from_lengths
-from fp16_optimizer import fp32_to_fp16, fp16_to_fp32
 
 
 class LocationLayer(nn.Module):
@@ -355,8 +354,6 @@ class Decoder(nn.Module):
             cell_input, (self.attention_hidden, self.attention_cell))
         self.attention_hidden = F.dropout(
             self.attention_hidden, self.p_attention_dropout, self.training)
-        self.attention_cell = F.dropout(
-            self.attention_cell, self.p_attention_dropout, self.training)
 
         attention_weights_cat = torch.cat(
             (self.attention_weights.unsqueeze(1),
@@ -372,8 +369,6 @@ class Decoder(nn.Module):
             decoder_input, (self.decoder_hidden, self.decoder_cell))
         self.decoder_hidden = F.dropout(
             self.decoder_hidden, self.p_decoder_dropout, self.training)
-        self.decoder_cell = F.dropout(
-            self.decoder_cell, self.p_decoder_dropout, self.training)
 
         decoder_hidden_attention_context = torch.cat(
             (self.decoder_hidden, self.attention_context), dim=1)
@@ -489,10 +484,6 @@ class Tacotron2(nn.Module):
             (text_padded, input_lengths, mel_padded, max_len, output_lengths),
             (mel_padded, gate_padded))
 
-    def parse_input(self, inputs):
-        inputs = fp32_to_fp16(inputs) if self.fp16_run else inputs
-        return inputs
-
     def parse_output(self, outputs, output_lengths=None):
         if self.mask_padding and output_lengths is not None:
             mask = ~get_mask_from_lengths(output_lengths)
@@ -503,20 +494,18 @@ class Tacotron2(nn.Module):
             outputs[1].data.masked_fill_(mask, 0.0)
             outputs[2].data.masked_fill_(mask[:, 0, :], 1e3)  # gate energies
 
-        outputs = fp16_to_fp32(outputs) if self.fp16_run else outputs
         return outputs
 
     def forward(self, inputs):
-        inputs, input_lengths, targets, max_len, \
-            output_lengths = self.parse_input(inputs)
-        input_lengths, output_lengths = input_lengths.data, output_lengths.data
+        text_inputs, text_lengths, mels, max_len, output_lengths = inputs
+        text_lengths, output_lengths = text_lengths.data, output_lengths.data
 
-        embedded_inputs = self.embedding(inputs).transpose(1, 2)
+        embedded_inputs = self.embedding(text_inputs).transpose(1, 2)
 
-        encoder_outputs = self.encoder(embedded_inputs, input_lengths)
+        encoder_outputs = self.encoder(embedded_inputs, text_lengths)
 
         mel_outputs, gate_outputs, alignments = self.decoder(
-            encoder_outputs, targets, memory_lengths=input_lengths)
+            encoder_outputs, mels, memory_lengths=text_lengths)
 
         mel_outputs_postnet = self.postnet(mel_outputs)
         mel_outputs_postnet = mel_outputs + mel_outputs_postnet
@@ -526,7 +515,6 @@ class Tacotron2(nn.Module):
             output_lengths)
 
     def inference(self, inputs):
-        inputs = self.parse_input(inputs)
         embedded_inputs = self.embedding(inputs).transpose(1, 2)
         encoder_outputs = self.encoder.inference(embedded_inputs)
         mel_outputs, gate_outputs, alignments = self.decoder.inference(

requirements.txt

@@ -1,6 +1,6 @@
 matplotlib==2.1.0
-tensorflow
-numpy==1.13.3
+tensorflow==1.1.0
+numpy==1.17.1
 inflect==0.2.5
 librosa==0.6.0
 scipy==1.0.0

stft.py

@@ -124,6 +124,7 @@ class STFT(torch.nn.Module):
                 np.where(window_sum > tiny(window_sum))[0])
             window_sum = torch.autograd.Variable(
                 torch.from_numpy(window_sum), requires_grad=False)
+            window_sum = window_sum.cuda() if magnitude.is_cuda else window_sum
             inverse_transform[:, :, approx_nonzero_indices] /= window_sum[approx_nonzero_indices]
 
             # scale by hop ratio

text_codec.py

@@ -0,0 +1,77 @@
+from utils import load_filepaths_and_text
+
+from text import text_to_sequence, sequence_to_text
+
+from hparams import create_hparams
+import sentencepiece as spm
+from text import symbols
+from bpemb import BPEmb
+
+
+SPM_CORPUS_FILE = "filelists/text_corpus.txt"
+SPM_MODEL_PREFIX = "spm"
+SPM_VOCAB_SIZE = 1000
+hparams = create_hparams()
+
+
+def _create_sentencepiece_corpus():
+    def get_text_list(text_file):
+        return [i[1] + "\n" for i in load_filepaths_and_text(text_file)]
+
+    full_text_list = get_text_list(hparams.training_files) + get_text_list(
+        hparams.validation_files
+    )
+    with open(SPM_CORPUS_FILE, "w") as fd:
+        fd.writelines(full_text_list)
+
+
+def _create_sentencepiece_vocab(vocab_size=SPM_VOCAB_SIZE):
+    train_params = "--input={} --model_type=unigram --character_coverage=1.0 --model_prefix={} --vocab_size={}".format(
+        SPM_CORPUS_FILE, SPM_MODEL_PREFIX, vocab_size
+    )
+    spm.SentencePieceTrainer.Train(train_params)
+
+
+def _spm_text_codecs():
+    sp = spm.SentencePieceProcessor()
+    sp.Load("{}.model".format(SPM_MODEL_PREFIX))
+
+    def ttseq(text, cleaners):
+        return sp.EncodeAsIds(text)
+
+    def seqtt(sequence):
+        return sp.DecodeIds(sequence)
+
+    return ttseq, seqtt
+
+
+def _bpemb_text_codecs():
+    bpemb_en = BPEmb(lang="en", dim=50, vs=148)
+    def ttseq(text, cleaners):
+        return bpemb_en.encode_ids(text)
+
+    def seqtt(sequence):
+        return bpemb_en.decode_ids(sequence)
+
+    return ttseq, seqtt
+
+# text_to_sequence, sequence_to_text = _spm_text_codecs()
+text_to_sequence, sequence_to_text = _bpemb_text_codecs()
+
+
+def _interactive_test():
+    prompt = "Hello world; how are you, doing ?"
+    while prompt not in ["q", "quit"]:
+        oup = sequence_to_text(text_to_sequence(prompt, hparams.text_cleaners))
+        print('==> ',oup)
+        prompt = input("> ")
+
+
+def main():
+    # _create_sentencepiece_corpus()
+    # _create_sentencepiece_vocab()
+    _interactive_test()
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -10,8 +10,6 @@ import torch.distributed as dist
 from torch.utils.data.distributed import DistributedSampler
 from torch.utils.data import DataLoader
 
-from fp16_optimizer import FP16_Optimizer
-
 from model import Tacotron2
 from data_utils import TextMelLoader, TextMelCollate
 from loss_function import Tacotron2Loss
@@ -19,19 +17,10 @@ from logger import Tacotron2Logger
 from hparams import create_hparams
 
 
-def batchnorm_to_float(module):
-    """Converts batch norm modules to FP32"""
-    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
-        module.float()
-    for child in module.children():
-        batchnorm_to_float(child)
-    return module
-
-
-def reduce_tensor(tensor, num_gpus):
+def reduce_tensor(tensor, n_gpus):
     rt = tensor.clone()
     dist.all_reduce(rt, op=dist.reduce_op.SUM)
-    rt /= num_gpus
+    rt /= n_gpus
     return rt
 
 
@@ -56,10 +45,14 @@ def prepare_dataloaders(hparams):
     valset = TextMelLoader(hparams.validation_files, hparams)
     collate_fn = TextMelCollate(hparams.n_frames_per_step)
 
-    train_sampler = DistributedSampler(trainset) \
-        if hparams.distributed_run else None
+    if hparams.distributed_run:
+        train_sampler = DistributedSampler(trainset)
+        shuffle = False
+    else:
+        train_sampler = None
+        shuffle = True
 
-    train_loader = DataLoader(trainset, num_workers=1, shuffle=False,
+    train_loader = DataLoader(trainset, num_workers=1, shuffle=shuffle,
                               sampler=train_sampler,
                               batch_size=hparams.batch_size, pin_memory=False,
                               drop_last=True, collate_fn=collate_fn)
@@ -80,22 +73,26 @@ def prepare_directories_and_logger(output_directory, log_directory, rank):
 def load_model(hparams):
     model = Tacotron2(hparams).cuda()
     if hparams.fp16_run:
-        model = batchnorm_to_float(model.half())
-        model.decoder.attention_layer.score_mask_value = float(finfo('float16').min)
+        model.decoder.attention_layer.score_mask_value = finfo('float16').min
 
     if hparams.distributed_run:
-        model = DistributedDataParallel(model)
-    elif torch.cuda.device_count() > 1:
-        model = DataParallel(model)
+        model = apply_gradient_allreduce(model)
 
     return model
 
 
-def warm_start_model(checkpoint_path, model):
+def warm_start_model(checkpoint_path, model, ignore_layers):
     assert os.path.isfile(checkpoint_path)
     print("Warm starting model from checkpoint '{}'".format(checkpoint_path))
     checkpoint_dict = torch.load(checkpoint_path, map_location='cpu')
-    model.load_state_dict(checkpoint_dict['state_dict'])
+    model_dict = checkpoint_dict['state_dict']
+    if len(ignore_layers) > 0:
+        model_dict = {k: v for k, v in model_dict.items()
+                      if k not in ignore_layers}
+        dummy_dict = model.state_dict()
+        dummy_dict.update(model_dict)
+        model_dict = dummy_dict
+    model.load_state_dict(model_dict)
     return model
 
 
@@ -137,15 +134,16 @@ def validate(model, criterion, valset, iteration, batch_size, n_gpus,
             y_pred = model(x)
             loss = criterion(y_pred, y)
             if distributed_run:
-                reduced_val_loss = reduce_tensor(loss.data, num_gpus).item()
+                reduced_val_loss = reduce_tensor(loss.data, n_gpus).item()
             else:
                 reduced_val_loss = loss.item()
             val_loss += reduced_val_loss
         val_loss = val_loss / (i + 1)
 
     model.train()
-    print("Validation loss {}: {:9f}  ".format(iteration, reduced_val_loss))
-    logger.log_validation(reduced_val_loss, model, y, y_pred, iteration)
+    if rank == 0:
+        print("Validation loss {}: {:9f}  ".format(iteration, reduced_val_loss))
+        logger.log_validation(reduced_val_loss, model, y, y_pred, iteration)
 
 
 def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
@@ -171,9 +169,11 @@ def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
     learning_rate = hparams.learning_rate
     optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate,
                                  weight_decay=hparams.weight_decay)
+
     if hparams.fp16_run:
-        optimizer = FP16_Optimizer(
-            optimizer, dynamic_loss_scale=hparams.dynamic_loss_scaling)
+        from apex import amp
+        model, optimizer = amp.initialize(
+            model, optimizer, opt_level='O2')
 
     if hparams.distributed_run:
         model = apply_gradient_allreduce(model)
@@ -190,7 +190,8 @@ def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
     epoch_offset = 0
     if checkpoint_path is not None:
         if warm_start:
-            model = warm_start_model(checkpoint_path, model)
+            model = warm_start_model(
+                checkpoint_path, model, hparams.ignore_layers)
         else:
             model, optimizer, _learning_rate, iteration = load_checkpoint(
                 checkpoint_path, model, optimizer)
@@ -200,6 +201,7 @@ def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
             epoch_offset = max(0, int(iteration / len(train_loader)))
 
     model.train()
+    is_overflow = False
     # ================ MAIN TRAINNIG LOOP! ===================
     for epoch in range(epoch_offset, hparams.epochs):
         print("Epoch: {}".format(epoch))
@@ -214,33 +216,36 @@ def train(output_directory, log_directory, checkpoint_path, warm_start, n_gpus,
 
             loss = criterion(y_pred, y)
             if hparams.distributed_run:
-                reduced_loss = reduce_tensor(loss.data, num_gpus).item()
+                reduced_loss = reduce_tensor(loss.data, n_gpus).item()
             else:
                 reduced_loss = loss.item()
-
             if hparams.fp16_run:
-                optimizer.backward(loss)
-                grad_norm = optimizer.clip_fp32_grads(hparams.grad_clip_thresh)
+                with amp.scale_loss(loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
             else:
                 loss.backward()
+
+            if hparams.fp16_run:
+                grad_norm = torch.nn.utils.clip_grad_norm_(
+                    amp.master_params(optimizer), hparams.grad_clip_thresh)
+                is_overflow = math.isnan(grad_norm)
+            else:
                 grad_norm = torch.nn.utils.clip_grad_norm_(
                     model.parameters(), hparams.grad_clip_thresh)
 
             optimizer.step()
 
-            overflow = optimizer.overflow if hparams.fp16_run else False
-
-            if not overflow and not math.isnan(reduced_loss) and rank == 0:
+            if not is_overflow and rank == 0:
                 duration = time.perf_counter() - start
                 print("Train loss {} {:.6f} Grad Norm {:.6f} {:.2f}s/it".format(
                     iteration, reduced_loss, grad_norm, duration))
                 logger.log_training(
                     reduced_loss, grad_norm, learning_rate, duration, iteration)
 
-            if not overflow and (iteration % hparams.iters_per_checkpoint == 0):
-                validate(model, criterion, valset, iteration, hparams.batch_size,
-                         n_gpus, collate_fn, logger, hparams.distributed_run, rank)
-
+            if not is_overflow and (iteration % hparams.iters_per_checkpoint == 0):
+                validate(model, criterion, valset, iteration,
+                         hparams.batch_size, n_gpus, collate_fn, logger,
+                         hparams.distributed_run, rank)
                 if rank == 0:
                     checkpoint_path = os.path.join(
                         output_directory, "checkpoint_{}".format(iteration))
@@ -259,7 +264,7 @@ if __name__ == '__main__':
     parser.add_argument('-c', '--checkpoint_path', type=str, default=None,
                         required=False, help='checkpoint path')
     parser.add_argument('--warm_start', action='store_true',
-                        help='load the model only (warm start)')
+                        help='load model weights only, ignore specified layers')
     parser.add_argument('--n_gpus', type=int, default=1,
                         required=False, help='number of gpus')
     parser.add_argument('--rank', type=int, default=0,