Introduce code from Rust monorepo. This is a mess

tts/nemo/core/config/__init__.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.core.config.base_config import Config
+from nemo.core.config.hydra_runner import hydra_runner
+from nemo.core.config.optimizers import (
+    AdadeltaParams,
+    AdagradParams,
+    AdamaxParams,
+    AdamParams,
+    AdamWParams,
+    NovogradParams,
+    OptimizerParams,
+    RMSpropParams,
+    RpropParams,
+    SGDParams,
+    get_optimizer_config,
+    register_optimizer_params,
+)
+from nemo.core.config.pytorch import DataLoaderConfig
+from nemo.core.config.pytorch_lightning import TrainerConfig
+from nemo.core.config.schedulers import (
+    CosineAnnealingParams,
+    InverseSquareRootAnnealingParams,
+    NoamAnnealingParams,
+    PolynomialDecayAnnealingParams,
+    PolynomialHoldDecayAnnealingParams,
+    SchedulerParams,
+    SquareAnnealingParams,
+    SquareRootAnnealingParams,
+    WarmupAnnealingParams,
+    WarmupHoldSchedulerParams,
+    WarmupSchedulerParams,
+    get_scheduler_config,
+    register_scheduler_params,
+)

tts/nemo/core/config/base_config.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Optional
+
+__all__ = ['Config']
+
+
+@dataclass
+class Config:
+    """
+    Abstract NeMo Configuration class.
+
+    Args:
+        name: name of the module/dataset/loss/model object (used in serialization, DEFAULT: None)
+    """
+
+    name: Optional[str] = None

tts/nemo/core/config/hydra_runner.py

@@ -0,0 +1,110 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import functools
+import os
+import sys
+from typing import Any, Callable, Optional
+
+from hydra._internal.utils import _run_hydra, get_args_parser
+from hydra.core.config_store import ConfigStore
+from hydra.types import TaskFunction
+from omegaconf import DictConfig
+
+
+def hydra_runner(
+    config_path: Optional[str] = ".", config_name: Optional[str] = None, schema: Optional[Any] = None
+) -> Callable[[TaskFunction], Any]:
+    """
+    Decorator used for passing the Config paths to main function.
+    Optionally registers a schema used for validation/providing default values.
+
+    Args:
+        config_path: Optional path that will be added to config search directory.
+            NOTE: The default value of `config_path` has changed between Hydra 1.0 and Hydra 1.1+.
+            Please refer to https://hydra.cc/docs/next/upgrades/1.0_to_1.1/changes_to_hydra_main_config_path/
+            for details.
+        config_name: Pathname of the config file.
+        schema: Structured config  type representing the schema used for validation/providing default values.
+    """
+
+    def decorator(task_function: TaskFunction) -> Callable[[], None]:
+        @functools.wraps(task_function)
+        def wrapper(cfg_passthrough: Optional[DictConfig] = None) -> Any:
+            # Check it config was passed.
+            if cfg_passthrough is not None:
+                return task_function(cfg_passthrough)
+            else:
+                args = get_args_parser()
+
+                # Parse arguments in order to retrieve overrides
+                parsed_args = args.parse_args()  # type: argparse.Namespace
+
+                # Get overriding args in dot string format
+                overrides = parsed_args.overrides  # type: list
+
+                # Disable the creation of .hydra subdir
+                # https://hydra.cc/docs/tutorials/basic/running_your_app/working_directory
+                overrides.append("hydra.output_subdir=null")
+                # Hydra logging outputs only to stdout (no log file).
+                # https://hydra.cc/docs/configure_hydra/logging
+                overrides.append("hydra/job_logging=stdout")
+
+                # Set run.dir ONLY for ExpManager "compatibility" - to be removed.
+                overrides.append("hydra.run.dir=.")
+
+                # Check if user set the schema.
+                if schema is not None:
+                    # Create config store.
+                    cs = ConfigStore.instance()
+
+                    # Get the correct ConfigStore "path name" to "inject" the schema.
+                    if parsed_args.config_name is not None:
+                        path, name = os.path.split(parsed_args.config_name)
+                        # Make sure the path is not set - as this will disable validation scheme.
+                        if path != '':
+                            sys.stderr.write(
+                                f"ERROR Cannot set config file path using `--config-name` when "
+                                "using schema. Please set path using `--config-path` and file name using "
+                                "`--config-name` separately.\n"
+                            )
+                            sys.exit(1)
+                    else:
+                        name = config_name
+
+                    # Register the configuration as a node under the name in the group.
+                    cs.store(name=name, node=schema)  # group=group,
+
+                # Wrap a callable object with name `parse_args`
+                # This is to mimic the ArgParser.parse_args() API.
+                class _argparse_wrapper:
+                    def __init__(self, arg_parser):
+                        self.arg_parser = arg_parser
+                        self._actions = arg_parser._actions
+
+                    def parse_args(self, args=None, namespace=None):
+                        return parsed_args
+
+                # no return value from run_hydra() as it may sometime actually run the task_function
+                # multiple times (--multirun)
+                _run_hydra(
+                    args_parser=_argparse_wrapper(args),
+                    task_function=task_function,
+                    config_path=config_path,
+                    config_name=config_name,
+                )
+
+        return wrapper
+
+    return decorator

tts/nemo/core/config/modelPT.py

@@ -0,0 +1,181 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Optional
+
+from omegaconf import MISSING
+
+from nemo.core import config
+from nemo.core.classes.dataset import DatasetConfig
+from nemo.utils import exp_manager
+
+
+@dataclass
+class SchedConfig:
+    name: str = MISSING
+    min_lr: float = 0.0
+    last_epoch: int = -1
+
+
+@dataclass
+class OptimConfig:
+    name: str = MISSING
+    sched: Optional[SchedConfig] = None
+
+
+@dataclass
+class ModelConfig:
+    """
+    Model component inside ModelPT
+    """
+
+    # ...
+    train_ds: Optional[DatasetConfig] = None
+    validation_ds: Optional[DatasetConfig] = None
+    test_ds: Optional[DatasetConfig] = None
+    optim: Optional[OptimConfig] = None
+
+
+@dataclass
+class HydraConfig:
+    run: Dict[str, Any] = field(default_factory=lambda: {"dir": "."})
+    job_logging: Dict[str, Any] = field(default_factory=lambda: {"root": {"handlers": None}})
+
+
+@dataclass
+class NemoConfig:
+    name: str = MISSING
+    model: ModelConfig = MISSING
+    trainer: config.TrainerConfig = config.TrainerConfig(
+        accelerator="ddp", checkpoint_callback=False, logger=False, log_every_n_steps=1
+    )
+    exp_manager: Optional[Any] = exp_manager.ExpManagerConfig()
+    hydra: HydraConfig = HydraConfig()
+
+
+class ModelConfigBuilder:
+    def __init__(self, model_cfg: ModelConfig):
+        """
+        Base class for any Model Config Builder.
+
+        A Model Config Builder is a utility class that accepts a ModelConfig dataclass,
+        and via a set of utility methods (that are implemented by the subclassed ModelConfigBuilder),
+        builds a finalized ModelConfig that can be supplied to a NemoModel dataclass as
+        the `model` component.
+
+        Subclasses *must* implement the private method `_finalize_cfg`.
+            Inside this method, they must update `self.model_cfg` with all interdependent config
+            options that need to be set (either updated by user explicitly or with their default value).
+
+            The updated model config must then be preserved in `self.model_cfg`.
+
+        Example:
+            # Create the config builder
+            config_builder = <subclass>ModelConfigBuilder()
+
+            # Update the components of the config that are modifiable
+            config_builder.set_X(X)
+            config_builder.set_Y(Y)
+
+            # Create a "finalized" config dataclass that will contain all the updates
+            # that were specified by the builder
+            model_config = config_builder.build()
+
+            # Use model config as is (or further update values), then create a new Model
+            model = nemo.<domain>.models.<ModelName>Model(cfg=model_config, trainer=Trainer())
+
+        Supported build methods:
+        -   set_train_ds: All model configs can accept a subclass of `DatasetConfig` as their
+                training config. Subclasses can override this method to enable auto-complete
+                by replacing `Optional[DatasetConfig]` with `Optional[<subclass of DatasetConfig>]`.
+
+        -   set_validation_ds: All model configs can accept a subclass of `DatasetConfig` as their
+                validation config. Subclasses can override this method to enable auto-complete
+                by replacing `Optional[DatasetConfig]` with `Optional[<subclass of DatasetConfig>]`.
+
+        -   set_test_ds: All model configs can accept a subclass of `DatasetConfig` as their
+                test config. Subclasses can override this method to enable auto-complete
+                by replacing `Optional[DatasetConfig]` with `Optional[<subclass of DatasetConfig>]`.
+
+        -   set_optim: A build method that supports changes to the Optimizer (and optionally,
+                the Scheduler) used for training the model. The function accepts two inputs -
+
+                `cfg`: A subclass of `OptimizerParams` - any OptimizerParams subclass can be used,
+                    in order to select an appropriate Optimizer. Examples: AdamParams.
+
+                `sched_cfg`: A subclass of `SchedulerParams` - any SchedulerParams subclass can be used,
+                    in order to select an appropriate Scheduler. Examples: CosineAnnealingParams.
+                    Note that this argument is optional.
+
+        -   build(): The method which should return a "finalized" ModelConfig dataclass.
+                Subclasses *should* always override this method, and update the signature
+                of this method with the return type of the Dataclass, so that it enables
+                autocomplete for the user.
+
+                Example:
+                    def build(self) -> EncDecCTCConfig:
+                        return super().build()
+
+        Any additional build methods must be added by subclasses of ModelConfigBuilder.
+
+        Args:
+            model_cfg:
+        """
+        self.model_cfg = model_cfg
+        self.train_ds_cfg = None
+        self.validation_ds_cfg = None
+        self.test_ds_cfg = None
+        self.optim_cfg = None
+
+    def set_train_ds(self, cfg: Optional[DatasetConfig] = None):
+        self.model_cfg.train_ds = cfg
+
+    def set_validation_ds(self, cfg: Optional[DatasetConfig] = None):
+        self.model_cfg.validation_ds = cfg
+
+    def set_test_ds(self, cfg: Optional[DatasetConfig] = None):
+        self.model_cfg.test_ds = cfg
+
+    def set_optim(self, cfg: config.OptimizerParams, sched_cfg: Optional[config.SchedulerParams] = None):
+        @dataclass
+        class WrappedOptimConfig(OptimConfig, cfg.__class__):
+            pass
+
+        # Setup optim
+        optim_name = cfg.__class__.__name__.replace("Params", "").lower()
+        wrapped_cfg = WrappedOptimConfig(name=optim_name, sched=None, **vars(cfg))
+
+        if sched_cfg is not None:
+
+            @dataclass
+            class WrappedSchedConfig(SchedConfig, sched_cfg.__class__):
+                pass
+
+            # Setup scheduler
+            sched_name = sched_cfg.__class__.__name__.replace("Params", "")
+            wrapped_sched_cfg = WrappedSchedConfig(name=sched_name, **vars(sched_cfg))
+
+            wrapped_cfg.sched = wrapped_sched_cfg
+
+        self.model_cfg.optim = wrapped_cfg
+
+    def _finalize_cfg(self):
+        raise NotImplementedError()
+
+    def build(self) -> ModelConfig:
+        # validate config
+        self._finalize_cfg()
+
+        return self.model_cfg

tts/nemo/core/config/optimizers.py

@@ -0,0 +1,265 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, Optional, Tuple
+
+from omegaconf import MISSING, OmegaConf
+
+__all__ = [
+    'OptimizerParams',
+    'AdamParams',
+    'NovogradParams',
+    'SGDParams',
+    'AdadeltaParams',
+    'AdamaxParams',
+    'AdagradParams',
+    'AdamWParams',
+    'RMSpropParams',
+    'RpropParams',
+]
+
+
+@dataclass
+class OptimizerParams:
+    """
+    Base Optimizer params with no values. User can chose it to explicitly override via
+    command line arguments
+    """
+
+    lr: Optional[float] = MISSING
+
+
+@dataclass
+class SGDParams(OptimizerParams):
+    """
+    Default configuration for Adam optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html?highlight=sgd#torch.optim.SGD
+    """
+
+    momentum: float = 0
+    dampening: float = 0
+    weight_decay: float = 0
+    nesterov: bool = False
+
+
+@dataclass
+class AdamParams(OptimizerParams):
+    """
+    Default configuration for Adam optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html?highlight=adam#torch.optim.Adam
+    """
+
+    # betas: Tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-08
+    weight_decay: float = 0
+    amsgrad: bool = False
+
+
+@dataclass
+class AdamWParams(OptimizerParams):
+    """
+    Default configuration for AdamW optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.AdamW
+    """
+
+    betas: Tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-08
+    weight_decay: float = 0
+    amsgrad: bool = False
+
+
+@dataclass
+class AdadeltaParams(OptimizerParams):
+    """
+    Default configuration for Adadelta optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.Adadelta
+    """
+
+    rho: float = 0.9
+    eps: float = 1e-6
+    weight_decay: float = 0
+
+
+@dataclass
+class AdamaxParams(OptimizerParams):
+    """
+    Default configuration for Adamax optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.Adamax
+    """
+
+    betas: Tuple[float, float] = (0.9, 0.999)
+    eps: float = 1e-8
+    weight_decay: float = 0
+
+
+@dataclass
+class AdagradParams(OptimizerParams):
+    """
+    Default configuration for Adagrad optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.Adagrad
+    """
+
+    lr_decay: float = 0
+    weight_decay: float = 0
+    initial_accumulator_value: float = 0
+    eps: float = 1e-10
+
+
+@dataclass
+class RMSpropParams(OptimizerParams):
+    """
+    Default configuration for RMSprop optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.RMSprop
+    """
+
+    alpha: float = 0.99
+    eps: float = 1e-8
+    weight_decay: float = 0
+    momentum: float = 0
+    centered: bool = False
+
+
+@dataclass
+class RpropParams(OptimizerParams):
+    """
+    Default configuration for RpropParams optimizer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/optim.html#torch.optim.Rprop
+    """
+
+    etas: Tuple[float, float] = (0.5, 1.2)
+    step_sizes: Tuple[float, float] = (1e-6, 50)
+
+
+@dataclass
+class NovogradParams(OptimizerParams):
+    """
+    Configuration of the Novograd optimizer.
+
+    It has been proposed  in "Stochastic Gradient Methods with Layer-wise
+    Adaptive Moments for Training of Deep Networks"
+    (https://arxiv.org/abs/1905.11286)
+
+    Args:
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper "On the Convergence of Adam and Beyond"
+    """
+
+    betas: Tuple[float, float] = (0.95, 0.98)
+    eps: float = 1e-8
+    weight_decay: float = 0
+    grad_averaging: bool = False
+    amsgrad: bool = False
+    luc: bool = False
+    luc_trust: float = 1e-3
+    luc_eps: float = 1e-8
+
+
+def register_optimizer_params(name: str, optimizer_params: OptimizerParams):
+    """
+    Checks if the optimizer param name exists in the registry, and if it doesnt, adds it.
+
+    This allows custom optimizer params to be added and called by name during instantiation.
+
+    Args:
+        name: Name of the optimizer. Will be used as key to retrieve the optimizer.
+        optimizer_params: Optimizer class
+    """
+    if name in AVAILABLE_OPTIMIZER_PARAMS:
+        raise ValueError(f"Cannot override pre-existing optimizers. Conflicting optimizer name = {name}")
+
+    AVAILABLE_OPTIMIZER_PARAMS[name] = optimizer_params
+
+
+def get_optimizer_config(name: str, **kwargs: Optional[Dict[str, Any]]) -> OptimizerParams:
+    """
+    Convenience method to obtain a OptimizerParams class and partially instantiate it with optimizer kwargs.
+
+    Args:
+        name: Name of the OptimizerParams in the registry.
+        kwargs: Optional kwargs of the optimizer used during instantiation.
+
+    Returns:
+        a partially instantiated OptimizerParams
+    """
+    if name is None:
+        return kwargs
+
+    if name not in AVAILABLE_OPTIMIZER_PARAMS:
+        raise ValueError(
+            f"Cannot resolve optimizer parameters '{name}'. Available optimizer parameters are : "
+            f"{AVAILABLE_OPTIMIZER_PARAMS.keys()}"
+        )
+
+    scheduler_params = AVAILABLE_OPTIMIZER_PARAMS[name]
+
+    if kwargs is not None and len(kwargs) != 0:
+        kwargs = OmegaConf.create(kwargs)
+        OmegaConf.merge(scheduler_params(), kwargs)
+
+    scheduler_params = partial(scheduler_params, **kwargs)
+    return scheduler_params
+
+
+AVAILABLE_OPTIMIZER_PARAMS = {
+    'optim_params': OptimizerParams,
+    'adam_params': AdamParams,
+    'novograd_params': NovogradParams,
+    'sgd_params': SGDParams,
+    'adadelta_params': AdadeltaParams,
+    'adamax_params': AdamaxParams,
+    'adagrad_params': AdagradParams,
+    'adamw_params': AdamWParams,
+    'rmsprop_params': RMSpropParams,
+    'rprop_params': RpropParams,
+}

tts/nemo/core/config/pytorch.py

@@ -0,0 +1,45 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from omegaconf import MISSING
+
+__all__ = ['DataLoaderConfig']
+
+
+@dataclass
+class DataLoaderConfig:
+    """
+    Configuration of PyTorch DataLoader.
+
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
+    """
+
+    batch_size: int = MISSING
+    shuffle: bool = False
+    sampler: Optional[Any] = None
+    batch_sampler: Optional[Any] = None
+    num_workers: int = 0
+    collate_fn: Optional[Any] = None
+    pin_memory: bool = False
+    drop_last: bool = False
+    timeout: int = 0
+    worker_init_fn: Optional[Any] = None
+    multiprocessing_context: Optional[Any] = None

tts/nemo/core/config/pytorch_lightning.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from hydra.core.config_store import ConfigStore
+
+__all__ = ['TrainerConfig']
+
+
+cs = ConfigStore.instance()
+
+
+@dataclass
+class TrainerConfig:
+    """
+    Configuration of PyTorch Lightning Trainer.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    ..warning:
+        Picked just few params of the PTL trainer for now. This needs to be discussed.
+    ..note:
+        For the details on the function/meanings of the arguments, please refer to:
+        https://pytorch-lightning.readthedocs.io/en/latest/trainer.html#
+    """
+
+    logger: Any = True
+    checkpoint_callback: Any = True
+    callbacks: Optional[Any] = None
+    default_root_dir: Optional[str] = None
+    gradient_clip_val: float = 0
+    process_position: int = 0
+    num_nodes: int = 1
+    num_processes: int = 1
+    gpus: Optional[Any] = None
+    auto_select_gpus: bool = False
+    tpu_cores: Optional[Any] = None
+    log_gpu_memory: Optional[str] = None
+    progress_bar_refresh_rate: int = 1
+    overfit_batches: Any = 0.0
+    track_grad_norm: Any = -1
+    check_val_every_n_epoch: int = 1
+    fast_dev_run: bool = False
+    accumulate_grad_batches: Any = 1
+    max_epochs: int = 1000
+    min_epochs: int = 1
+    max_steps: Optional[int] = None
+    min_steps: Optional[int] = None
+    limit_train_batches: Any = 1.0
+    limit_val_batches: Any = 1.0
+    limit_test_batches: Any = 1.0
+    val_check_interval: Any = 1.0
+    flush_logs_every_n_steps: int = 100
+    log_every_n_steps: int = 50
+    accelerator: Optional[str] = None
+    sync_batchnorm: bool = False
+    precision: int = 32
+    weights_summary: Optional[str] = "full"  # ModelSummary.MODE_DEFAULT
+    weights_save_path: Optional[str] = None
+    num_sanity_val_steps: int = 2
+    truncated_bptt_steps: Optional[int] = None
+    resume_from_checkpoint: Optional[str] = None
+    profiler: Optional[Any] = None
+    benchmark: bool = False
+    deterministic: bool = False
+    reload_dataloaders_every_epoch: bool = False
+    auto_lr_find: Any = False
+    replace_sampler_ddp: bool = True
+    terminate_on_nan: bool = False
+    auto_scale_batch_size: Any = False
+    prepare_data_per_node: bool = True
+    amp_backend: str = 'native'
+    amp_level: str = 'O2'  # backward compatible, todo: remove in v1.0.0
+    plugins: Optional[Any] = None  # Optional[Union[str, list]]
+    move_metrics_to_cpu: bool = False
+    multiple_trainloader_mode: str = 'max_size_cycle'
+    limit_predict_batches: float = 1.0
+    stochastic_weight_avg: bool = False
+    gradient_clip_algorithm: str = 'norm'
+    max_time: Optional[Any] = None  # can be one of Union[str, timedelta, Dict[str, int], None]
+
+
+# Register the trainer config.
+cs.store(
+    group="trainer", name="trainer", node=TrainerConfig,
+)

tts/nemo/core/config/schedulers.py

@@ -0,0 +1,251 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class SchedulerParams:
+    """
+    Base configuration for all schedulers.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    last_epoch: int = -1
+
+
+@dataclass
+class WarmupSchedulerParams(SchedulerParams):
+    """
+    Base configuration for all schedulers.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    max_steps: int = 0
+    warmup_steps: Optional[float] = None
+    warmup_ratio: Optional[float] = None
+
+
+@dataclass
+class WarmupHoldSchedulerParams(WarmupSchedulerParams):
+    """
+    Base configuration for all schedulers.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    hold_steps: Optional[float] = None
+    hold_ratio: Optional[float] = None
+    min_lr: float = 0.0
+
+
+@dataclass
+class SquareAnnealingParams(WarmupSchedulerParams):
+    """
+    Square Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    min_lr: float = 1e-5
+
+
+@dataclass
+class SquareRootAnnealingParams(WarmupSchedulerParams):
+    """
+    Square Root Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    min_lr: float = 0.0
+
+
+@dataclass
+class CosineAnnealingParams(WarmupSchedulerParams):
+    """
+    Cosine Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    min_lr: float = 0.0
+
+
+@dataclass
+class NoamAnnealingParams(WarmupSchedulerParams):
+    """
+    Cosine Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    min_lr: float = 0.0
+
+
+@dataclass
+class WarmupAnnealingParams(WarmupSchedulerParams):
+    """
+    Warmup Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    warmup_ratio: 0.0
+
+
+@dataclass
+class InverseSquareRootAnnealingParams(WarmupSchedulerParams):
+    """
+    Inverse Square Root Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+
+@dataclass
+class PolynomialDecayAnnealingParams(WarmupSchedulerParams):
+    """
+    Polynomial Decay Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    power: float = 1.0
+    cycle: bool = False
+
+
+@dataclass
+class PolynomialHoldDecayAnnealingParams(WarmupSchedulerParams):
+    """
+    Polynomial Hold Decay Annealing parameter config
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    power: float = 1.0
+    cycle: bool = False
+
+
+"""
+Pytorch Optimizers
+"""
+
+
+@dataclass
+class StepLRParams(SchedulerParams):
+    """
+    Config for StepLR.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    step_size: float = 0.1
+    gamma: float = 0.1
+
+
+@dataclass
+class ExponentialLRParams(SchedulerParams):
+    """
+    Config for ExponentialLR.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    gamma: float = 0.9
+
+
+@dataclass
+class ReduceLROnPlateauParams:
+    """
+    Config for ReduceLROnPlateau.
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    mode: str = 'min'
+    factor: float = 0.1
+    patience: int = 10
+    verbose: bool = False
+    threshold: float = 1e-4
+    threshold_mode: str = 'rel'
+    cooldown: int = 0
+    min_lr: float = 0
+    eps: float = 1e-8
+
+
+@dataclass
+class CyclicLRParams(SchedulerParams):
+    """
+    Config for CyclicLR.
+    NOTE:
+    # `scale_fn` is not supported
+
+    It is not derived from Config as it is not a NeMo object (and in particular it doesn't need a name).
+    """
+
+    base_lr: float = 0.001
+    max_lr: float = 0.1
+    step_size_up: int = 2000
+    step_size_down: Optional[int] = None
+    mode: str = 'triangular'
+    gamma: float = 1.0
+    scale_mode: str = 'cycle'
+    # scale_fn is not supported
+    cycle_momentum: bool = True
+    base_momentum: float = 0.8
+    max_momentum: float = 0.9
+
+
+def register_scheduler_params(name: str, scheduler_params: SchedulerParams):
+    """
+    Checks if the schduler config name exists in the registry, and if it doesnt, adds it.
+
+    This allows custom schedulers to be added and called by name during instantiation.
+
+    Args:
+        name: Name of the optimizer. Will be used as key to retrieve the optimizer.
+        scheduler_params: SchedulerParams class
+    """
+    if name in AVAILABLE_SCHEDULER_PARAMS:
+        raise ValueError(f"Cannot override pre-existing optimizers. Conflicting optimizer name = {name}")
+
+    AVAILABLE_SCHEDULER_PARAMS[name] = scheduler_params
+
+
+def get_scheduler_config(name: str, **kwargs: Optional[Dict[str, Any]]) -> SchedulerParams:
+    """
+    Convenience method to obtain a SchedulerParams class and partially instantiate it with optimizer kwargs.
+
+    Args:
+        name: Name of the SchedulerParams in the registry.
+        kwargs: Optional kwargs of the optimizer used during instantiation.
+
+    Returns:
+        a partially instantiated SchedulerParams
+    """
+    if name not in AVAILABLE_SCHEDULER_PARAMS:
+        raise ValueError(
+            f"Cannot resolve scheduler parameters '{name}'. Available scheduler parameters are : "
+            f"{AVAILABLE_SCHEDULER_PARAMS.keys()}"
+        )
+
+    scheduler_params = AVAILABLE_SCHEDULER_PARAMS[name]
+    scheduler_params = partial(scheduler_params, **kwargs)
+    return scheduler_params
+
+
+AVAILABLE_SCHEDULER_PARAMS = {
+    'SchedulerParams': SchedulerParams,
+    'WarmupPolicyParams': WarmupSchedulerParams,
+    'WarmupHoldPolicyParams': WarmupHoldSchedulerParams,
+    'SquareAnnealingParams': SquareAnnealingParams,
+    'SquareRootAnnealingParams': SquareRootAnnealingParams,
+    'InverseSquareRootAnnealingParams': InverseSquareRootAnnealingParams,
+    'CosineAnnealingParams': CosineAnnealingParams,
+    'NoamAnnealingParams': NoamAnnealingParams,
+    'WarmupAnnealingParams': WarmupAnnealingParams,
+    'PolynomialDecayAnnealingParams': PolynomialDecayAnnealingParams,
+    'PolynomialHoldDecayAnnealingParams': PolynomialHoldDecayAnnealingParams,
+}