Add pretrained wavernn

Author: yangarbiter

torchaudio/models/__init__.py

@@ -1,5 +1,5 @@
 from .wav2letter import Wav2Letter
-from .wavernn import WaveRNN
+from .wavernn import WaveRNN, wavernn_10k_epochs_8bits_ljspeech
 from .conv_tasnet import ConvTasNet
 from .deepspeech import DeepSpeech
 from .wav2vec2 import (
@@ -13,6 +13,7 @@
 __all__ = [
     'Wav2Letter',
     'WaveRNN',
+    'wavernn_10k_epochs_8bits_ljspeech',
     'ConvTasNet',
     'DeepSpeech',
     'Wav2Vec2Model',

torchaudio/models/_utils.py

@@ -0,0 +1,8 @@
+try:
+    from torch.hub import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+__all__ = [
+    'load_state_dict_from_url',
+]

torchaudio/models/wavernn.py

@@ -1,18 +1,28 @@
-from typing import List, Tuple
+from typing import List, Tuple, Any
 
 import torch
 from torch import Tensor
 from torch import nn
 
+from ._utils import load_state_dict_from_url
+
+
 __all__ = [
     "ResBlock",
     "MelResNet",
     "Stretch2d",
     "UpsampleNetwork",
     "WaveRNN",
+    "wavernn_10k_epochs_8bits_ljspeech",
 ]
 
 
+model_urls = {
+    'wavernn_10k_epochs_8bits_ljspeech': 'https://download.pytorch.org/models/'
+                                         'audio/wavernn_10k_epochs_8bits_ljspeech.pth',
+}
+
+
 class ResBlock(nn.Module):
     r"""ResNet block based on *Efficient Neural Audio Synthesis* [:footcite:`kalchbrenner2018efficient`].
 
@@ -324,3 +334,37 @@ def forward(self, waveform: Tensor, specgram: Tensor) -> Tensor:
 
         # bring back channel dimension
         return x.unsqueeze(1)
+
+
+def _wavernn(arch: str, pretrained: bool, progress: bool, **kwargs: Any) -> WaveRNN:
+    model = WaveRNN(**kwargs)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls['wavernn'],
+                                              progress=progress)
+        model.load_state_dict(state_dict)
+    return model
+
+
+def wavernn_10k_epochs_8bits_ljspeech(pretrained: bool = True, progress: bool = True, **kwargs: Any) -> WaveRNN:
+    r"""WaveRNN model trained with 10k epochs and 8 bits depth waveform on the LJSpeech dataset.
+    The model is trained using the default parameters and code of the examples/pipeline_wavernn/main.py.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on LJSpeech
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    n_bits = 8
+    configs = {
+        'upsample_scales': [5, 5, 11],
+        'n_classes': 2 ** n_bits,
+        'hop_length': 275,
+        'n_res_block': 10,
+        'n_rnn': 512,
+        'n_fc': 512,
+        'kernel_size': 5,
+        'n_freq': 80,
+        'n_hidden': 128,
+        'n_output': 128
+    }
+    configs.update(kwargs)
+    return _wavernn("wavernn_10k_epochs_8bits_ljspeech", pretrained=pretrained, progress=progress, **configs)