Move Tacotron2 out of prototype (#1714)

pytorch · Aug 18, 2021 · 352d63c · 352d63c
1 parent 0f603eb
commit 352d63c
Show file tree

Hide file tree

Showing 8 changed files with 27 additions and 43 deletions.
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -40,7 +40,6 @@ The :mod:`torchaudio` package consists of I/O, popular datasets and common audio
    kaldi_io
    utils
    rnnt_loss
-   tacotron2
 
 
 .. toctree::

diff --git a/docs/source/models.rst b/docs/source/models.rst
@@ -25,6 +25,24 @@ DeepSpeech
   .. automethod:: forward
 
 
+Tacotron2
+~~~~~~~~~
+
+.. autoclass:: Tacotron2
+
+  .. automethod:: forward
+
+  .. automethod:: infer
+
+Factory Functions
+-----------------
+
+tacotron2
+---------
+
+.. autofunction:: tacotron2
+
+
 Wav2Letter
 ~~~~~~~~~~
 

diff --git a/docs/source/tacotron2.rst b/docs/source/tacotron2.rst
diff --git a/examples/pipeline_tacotron2/inference.py b/examples/pipeline_tacotron2/inference.py
@@ -11,8 +11,8 @@
 import torch
 import torchaudio
 import numpy as np
-from torchaudio.prototype.tacotron2 import Tacotron2
-from torchaudio.prototype.tacotron2 import tacotron2 as pretrained_tacotron2
+from torchaudio.models import Tacotron2
+from torchaudio.models import tacotron2 as pretrained_tacotron2
 
 from utils import prepare_input_sequence
 from datasets import InverseSpectralNormalization
@@ -28,7 +28,7 @@ def parse_args():
     r"""
     Parse commandline arguments.
     """
-    from torchaudio.prototype.tacotron2 import _MODEL_CONFIG_AND_URLS as tacotron2_config_and_urls
+    from torchaudio.models.tacotron2 import _MODEL_CONFIG_AND_URLS as tacotron2_config_and_urls
     from torchaudio.models.wavernn import _MODEL_CONFIG_AND_URLS as wavernn_config_and_urls
 
     parser = argparse.ArgumentParser(description=__doc__)

diff --git a/examples/pipeline_tacotron2/train.py b/examples/pipeline_tacotron2/train.py
@@ -44,7 +44,7 @@
 from torch.utils.tensorboard import SummaryWriter
 from torch.utils.data import DataLoader
 from torch.optim import Adam
-from torchaudio.prototype.tacotron2 import Tacotron2
+from torchaudio.models import Tacotron2
 from tqdm import tqdm
 import matplotlib.pyplot as plt
 plt.switch_backend('agg')

diff --git a/test/torchaudio_unittest/models/tacotron2/model_test_impl.py b/test/torchaudio_unittest/models/tacotron2/model_test_impl.py
@@ -1,7 +1,8 @@
 from typing import Tuple
 import torch
 from torch import Tensor
-from torchaudio.prototype.tacotron2 import Tacotron2, _Encoder, _Decoder
+from torchaudio.models import Tacotron2
+from torchaudio.models.tacotron2 import _Encoder, _Decoder
 from torchaudio_unittest.common_utils import (
     TestBaseMixin,
     TempDirMixin,

diff --git a/torchaudio/models/__init__.py b/torchaudio/models/__init__.py
@@ -2,6 +2,7 @@
 from .wavernn import WaveRNN, wavernn
 from .conv_tasnet import ConvTasNet
 from .deepspeech import DeepSpeech
+from .tacotron2 import Tacotron2, tacotron2
 from .wav2vec2 import (
     Wav2Vec2Model,
     wav2vec2_base,
@@ -20,4 +21,6 @@
     'wav2vec2_base',
     'wav2vec2_large',
     'wav2vec2_large_lv60k',
+    'Tacotron2',
+    'tacotron2',
 ]
diff --git a/torchaudio/prototype/tacotron2.py → torchaudio/models/tacotron2.py b/torchaudio/prototype/tacotron2.py → torchaudio/models/tacotron2.py