Introduce code from Rust monorepo. This is a mess

tts/nemo/core/__init__.py

@@ -0,0 +1,16 @@
+# 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 nemo.core.neural_types
+from nemo.core.classes import *

tts/nemo/core/classes/__init__.py

@@ -0,0 +1,29 @@
+# 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.classes.common import (
+    FileIO,
+    Model,
+    PretrainedModelInfo,
+    Serialization,
+    Typing,
+    is_typecheck_enabled,
+    typecheck,
+)
+from nemo.core.classes.dataset import Dataset, IterableDataset
+from nemo.core.classes.exportable import Exportable, ExportFormat
+from nemo.core.classes.loss import Loss
+from nemo.core.classes.modelPT import ModelPT
+from nemo.core.classes.module import NeuralModule

tts/nemo/core/classes/common.py

@@ -0,0 +1,812 @@
+# 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.
+
+
+"""Interfaces common to all Neural Modules and Models."""
+import hashlib
+import traceback
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from enum import Enum
+from functools import total_ordering
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import hydra
+import wrapt
+from omegaconf import DictConfig, OmegaConf
+
+import nemo
+from nemo.core.neural_types import NeuralType, NeuralTypeComparisonResult
+from nemo.utils import logging
+from nemo.utils.cloud import maybe_download_from_cloud
+from nemo.utils.model_utils import import_class_by_path, maybe_update_config_version
+
+__all__ = ['Typing', 'FileIO', 'Model', 'Serialization', 'typecheck']
+
+
+_TYPECHECK_ENABLED = True
+
+
+def is_typecheck_enabled():
+    """
+    Getter method for typechecking state.
+    """
+    return _TYPECHECK_ENABLED
+
+
+@dataclass
+class TypecheckMetadata:
+    """
+    Metadata class for input/output neural types.
+
+    # Primary attributes
+    original_types: Preserve the dictionary of type information provided.
+
+    ignore_collections: For backward compatibility, container support can be disabled explicitly
+        using this flag. When set to True, all nesting is ignored and nest-depth checks are skipped.
+
+    # Derived attributed
+    mandatory_types: Sub-dictionary of `original_types` which contains only those types which
+        are mandatory to include when calling the function.
+
+    base_types: Dictionary of flattened `str: NeuralType` definitions, disregarding the nest level
+        details into appropriate arguments.
+
+    container_depth: Dictionary mapping `str: int` - such that the valid depth of the nest of this
+        neural type is recorded.
+
+    has_container_types: Bool flag declaring if any of the neural types declares a container nest
+        in its signature.
+
+    is_singular_container_type: Bool flag declaring if this is a single Neural Type with a container
+        nest in its signature. Required for supporting python list expansion in return statement.
+
+    """
+
+    original_types: Dict[str, NeuralType]
+    ignore_collections: bool
+
+    mandatory_types: Dict[str, NeuralType] = field(init=False)
+    base_types: Dict[str, NeuralType] = field(init=False)
+
+    container_depth: Dict[str, int] = field(init=False)
+    has_container_types: bool = field(init=False)
+    is_singular_container_type: bool = field(init=False)
+
+    def __post_init__(self):
+        # If even one NeuralType declares a container nest, set to True
+        has_container_types = False
+        for type_val in self.original_types.values():
+            if isinstance(type_val, (list, tuple)):
+                has_container_types = True
+                break
+        self.has_container_types = has_container_types
+
+        # If only one NeuralType is declared, and it declares a container nest, set to True
+        if self.has_container_types and len(self.original_types) == 1:
+            self.is_singular_container_type = True
+        else:
+            self.is_singular_container_type = False
+
+        # If container nests are declared, flatten the nest into `base_types`
+        # Also compute the nest depth for each of the NeuralTypes
+        if self.has_container_types:
+            self.base_types = {}
+            self.container_depth = {}
+
+            for type_key, type_val in self.original_types.items():
+                depth = 0
+                while isinstance(type_val, (list, tuple)):
+                    if len(type_val) > 1:
+                        raise TypeError(
+                            f"Neural Type `{type_key}`: {type_val} definition contains more than one element when"
+                            "declaring the nested container structure.\n"
+                            "Please ensure that you have only 1 NeuralType inside of the entire nested structure "
+                            "definition."
+                        )
+
+                    type_val = type_val[0]
+                    depth += 1
+
+                self.base_types[type_key] = type_val
+                self.container_depth[type_key] = depth
+        else:
+            # Otherwise, simply preserve the original_types and set depth of nest to 0.
+            self.base_types = self.original_types
+            self.container_depth = {type_key: 0 for type_key in self.base_types.keys()}
+
+        # Compute subset of original_types which are mandatory in the call argspec
+        self.mandatory_types = {
+            type_key: type_val for type_key, type_val in self.base_types.items() if not type_val.optional
+        }
+
+
+class Typing(ABC):
+    """
+    An interface which endows module with neural types
+    """
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Define these to enable input neural type checks"""
+        return None
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """Define these to enable output neural type checks"""
+        return None
+
+    def _validate_input_types(self, input_types=None, ignore_collections=False, **kwargs):
+        """
+        This function does a few things.
+        1) It ensures that len(self.input_types <non-optional>) <= len(kwargs) <= len(self.input_types).
+        2) For each (keyword name, keyword value) passed as input to the wrapped function:
+            - Check if the keyword name exists in the list of valid self.input_types names.
+            - Check if keyword value has the `neural_type` property.
+                - If it does, then perform a comparative check and assert that neural types
+                    are compatible (SAME or GREATER).
+            - Check if keyword value is a container type (list or tuple). If yes,
+                then perform the elementwise test of neural type above on each element
+                of the nested structure, recursively.
+
+        Args:
+            input_types: Either the `input_types` defined at class level, or the local function
+                overridden type definition.
+            ignore_collections: For backward compatibility, container support can be disabled explicitly
+                using this flag. When set to True, all nesting is ignored and nest-depth checks are skipped.
+            kwargs: Dictionary of argument_name:argument_value pairs passed to the wrapped
+                function upon call.
+        """
+        # TODO: Properly implement this
+        if input_types is not None:
+            # Precompute metadata
+            metadata = TypecheckMetadata(original_types=input_types, ignore_collections=ignore_collections)
+
+            total_input_types = len(input_types)
+            mandatory_input_types = len(metadata.mandatory_types)
+
+            # Allow number of input arguments to be <= total input neural types.
+            if len(kwargs) < mandatory_input_types or len(kwargs) > total_input_types:
+                raise TypeError(
+                    f"Number of input arguments provided ({len(kwargs)}) is not as expected. Function has "
+                    f"{total_input_types} total inputs with {mandatory_input_types} mandatory inputs."
+                )
+
+            for key, value in kwargs.items():
+                # Check if keys exists in the defined input types
+                if key not in input_types:
+                    raise TypeError(
+                        f"Input argument {key} has no corresponding input_type match. "
+                        f"Existing input_types = {input_types.keys()}"
+                    )
+
+                # Perform neural type check
+                if hasattr(value, 'neural_type') and not metadata.base_types[key].compare(value.neural_type) in (
+                    NeuralTypeComparisonResult.SAME,
+                    NeuralTypeComparisonResult.GREATER,
+                ):
+                    error_msg = [
+                        f"{input_types[key].compare(value.neural_type)} :",
+                        f"Input type expected : {input_types[key]}",
+                        f"Input type found : {value.neural_type}",
+                        f"Argument: {key}",
+                    ]
+                    for i, dict_tuple in enumerate(metadata.base_types[key].elements_type.type_parameters.items()):
+                        error_msg.insert(i + 2, f'  input param_{i} : {dict_tuple[0]}: {dict_tuple[1]}')
+                    for i, dict_tuple in enumerate(value.neural_type.elements_type.type_parameters.items()):
+                        error_msg.append(f'  input param_{i} : {dict_tuple[0]}: {dict_tuple[1]}')
+                    raise TypeError("\n".join(error_msg))
+
+                # Perform input ndim check
+                if hasattr(value, 'shape'):
+                    value_shape = value.shape
+                    type_shape = metadata.base_types[key].axes
+                    name = key
+
+                    if type_shape is not None and len(value_shape) != len(type_shape):
+                        raise TypeError(
+                            f"Input shape mismatch occured for {name} in module {self.__class__.__name__} : \n"
+                            f"Input shape expected = {metadata.base_types[key].axes} | \n"
+                            f"Input shape found : {value_shape}"
+                        )
+
+                # Perform recursive neural type check for homogeneous elements
+                elif isinstance(value, list) or isinstance(value, tuple):
+                    for ind, val in enumerate(value):
+                        """
+                        This initiates a DFS, tracking the depth count as it goes along the nested structure.
+                        Initial depth is 1 as we consider the current loop to be the 1st step inside the nest.
+                        """
+                        self.__check_neural_type(val, metadata, depth=1, name=key)
+
+    def _attach_and_validate_output_types(self, out_objects, ignore_collections=False, output_types=None):
+        """
+        This function does a few things.
+        1) It ensures that len(out_object) == len(self.output_types).
+        2) If the output is a tensor (or list/tuple of list/tuple ... of tensors), it
+            attaches a neural_type to it. For objects without the neural_type attribute,
+            such as python objects (dictionaries and lists, primitive data types, structs),
+            no neural_type is attached.
+
+            Note: tensor.neural_type is only checked during _validate_input_types which is
+            called prior to forward().
+
+        Args:
+            output_types: Either the `output_types` defined at class level, or the local function
+                overridden type definition.
+            ignore_collections: For backward compatibility, container support can be disabled explicitly
+                using this flag. When set to True, all nesting is ignored and nest-depth checks are skipped.
+            out_objects: The outputs of the wrapped function.
+        """
+        # TODO: Properly implement this
+        if output_types is not None:
+            # Precompute metadata
+            metadata = TypecheckMetadata(original_types=output_types, ignore_collections=ignore_collections)
+            out_types_list = list(metadata.base_types.items())
+
+            # First convert all outputs to list/tuple format to check correct number of outputs
+            if type(out_objects) in (list, tuple):
+                out_container = out_objects  # can be any rank nested structure
+            else:
+                out_container = [out_objects]
+
+            # If this neural type has a *single output*, with *support for nested outputs*,
+            # then *do not* perform any check on the number of output items against the number
+            # of neural types (in this case, 1).
+            # This is done as python will *not* wrap a single returned list into a tuple of length 1,
+            # instead opting to keep the list intact. Therefore len(out_container) in such a case
+            # is the length of all the elements of that list - each of which has the same corresponding
+            # neural type (defined as the singular container type).
+            if metadata.is_singular_container_type:
+                pass
+
+            # In all other cases, python will wrap multiple outputs into an outer tuple.
+            # As such, now the number of elements in this outer tuple should exactly match
+            # the number of output types defined.
+            elif len(out_types_list) != len(out_container):
+                raise TypeError(
+                    "Number of output arguments provided ({}) is not as expected ({}).\n"
+                    "This can be either because insufficient number of output NeuralTypes were provided,"
+                    "or the provided NeuralTypes {} should enable container support "
+                    "(add '[]' to the NeuralType definition)".format(
+                        len(out_container), len(output_types), output_types
+                    )
+                )
+
+            # Attach types recursively, if possible
+            if not isinstance(out_objects, tuple) and not isinstance(out_objects, list):
+                # Here, out_objects is a single object which can potentially be attached with a NeuralType
+                try:
+                    out_objects.neural_type = out_types_list[0][1]
+                except Exception:
+                    pass
+
+                # Perform output ndim check
+                if hasattr(out_objects, 'shape'):
+                    value_shape = out_objects.shape
+                    type_shape = out_types_list[0][1].axes
+                    name = out_types_list[0][0]
+
+                    if type_shape is not None and len(value_shape) != len(type_shape):
+                        raise TypeError(
+                            f"Output shape mismatch occured for {name} in module {self.__class__.__name__} : \n"
+                            f"Output shape expected = {type_shape} | \n"
+                            f"Output shape found : {value_shape}"
+                        )
+
+            elif metadata.is_singular_container_type:
+                # If only a single neural type is provided, and it defines a container nest,
+                # then all elements of the returned list/tuple are assumed to belong to that
+                # singular neural type.
+                # As such, the "current" depth inside the DFS loop is counted as 1,
+                # and subsequent nesting will increase this count.
+
+                # NOTE:
+                # As the flag `is_singular_container_type` will activate only for
+                # the case where there is 1 output type defined with container nesting,
+                # this is a safe assumption to make.
+                depth = 1
+
+                # NOTE:
+                # A user may chose to explicitly wrap the single output list within an explicit tuple
+                # In such a case we reduce the "current" depth to 0 - to acknowledge the fact that
+                # the actual nest exists within a wrapper tuple.
+                if len(out_objects) == 1 and type(out_objects) == tuple:
+                    depth = 0
+
+                for ind, res in enumerate(out_objects):
+                    self.__attach_neural_type(res, metadata, depth=depth, name=out_types_list[0][0])
+            else:
+                # If more then one item is returned in a return statement, python will wrap
+                # the output with an outer tuple. Therefore there must be a 1:1 correspondence
+                # of the output_neural type (with or without nested structure) to the actual output
+                # (whether it is a single object or a nested structure of objects).
+                # Therefore in such a case, we "start" the DFS at depth 0 - since the recursion is
+                # being applied on 1 neural type : 1 output struct (single or nested output).
+                # Since we are guarenteed that the outer tuple will be built by python,
+                # assuming initial depth of 0 is appropriate.
+                for ind, res in enumerate(out_objects):
+                    self.__attach_neural_type(res, metadata, depth=0, name=out_types_list[ind][0])
+
+    def __check_neural_type(self, obj, metadata, depth, name=None):
+        if isinstance(obj, tuple) or isinstance(obj, list):
+            for elem in obj:
+                self.__check_neural_type(elem, metadata, depth + 1, name=name)
+            return  # after processing nest, return to avoid testing nest itself
+
+        type_val = metadata.base_types[name]
+
+        # If nest depth doesnt match neural type structure depth, raise an error
+        if not metadata.ignore_collections and depth != metadata.container_depth[name]:
+            raise TypeError(
+                "While checking input neural types,\n"
+                "Nested depth of value did not match container specification:\n"
+                f"Current nested depth of NeuralType '{name}' ({type_val}): {depth}\n"
+                f"Expected nested depth : {metadata.container_depth[name]}"
+            )
+
+        if hasattr(obj, 'neural_type') and not type_val.compare(obj.neural_type) in (
+            NeuralTypeComparisonResult.SAME,
+            NeuralTypeComparisonResult.GREATER,
+        ):
+            raise TypeError(
+                f"{type_val.compare(obj.neural_type)} : \n"
+                f"Input type expected = {type_val} | \n"
+                f"Input type found : {obj.neural_type}"
+            )
+
+        # Perform input ndim check
+        if hasattr(obj, 'shape'):
+            value_shape = obj.shape
+            type_shape = type_val.axes
+
+            if type_shape is not None and len(value_shape) != len(type_shape):
+                raise TypeError(
+                    f"Input shape mismatch occured for {name} in module {self.__class__.__name__} : \n"
+                    f"Input shape expected = {type_shape} | \n"
+                    f"Input shape found : {value_shape}"
+                )
+
+    def __attach_neural_type(self, obj, metadata, depth, name=None):
+        if isinstance(obj, tuple) or isinstance(obj, list):
+            for elem in obj:
+                self.__attach_neural_type(elem, metadata, depth=depth + 1, name=name)
+            return  # after processing nest, return to avoid argument insertion into nest itself
+
+        type_val = metadata.base_types[name]
+
+        # If nest depth doesnt match neural type structure depth, raise an error
+        if not metadata.ignore_collections and depth != metadata.container_depth[name]:
+            raise TypeError(
+                "While attaching output neural types,\n"
+                "Nested depth of value did not match container specification:\n"
+                f"Current nested depth of NeuralType '{name}' ({type_val}): {depth}\n"
+                f"Expected nested depth : {metadata.container_depth[name]}"
+            )
+
+        try:
+            obj.neural_type = type_val
+        except Exception:
+            pass
+
+        # Perform output ndim check
+        if hasattr(obj, 'shape'):
+            value_shape = obj.shape
+            type_shape = type_val.axes
+
+            if type_shape is not None and len(value_shape) != len(type_shape):
+                raise TypeError(
+                    f"Output shape mismatch occured for {name} in module {self.__class__.__name__} : \n"
+                    f"Output shape expected = {type_shape} | \n"
+                    f"Output shape found : {value_shape}"
+                )
+
+
+class Serialization(ABC):
+    @classmethod
+    def from_config_dict(cls, config: DictConfig):
+        """Instantiates object using DictConfig-based configuration"""
+        # Resolve the config dict
+        if isinstance(config, DictConfig):
+            config = OmegaConf.to_container(config, resolve=True)
+            config = OmegaConf.create(config)
+            OmegaConf.set_struct(config, True)
+
+        config = maybe_update_config_version(config)
+
+        # Hydra 0.x API
+        if ('cls' in config or 'target' in config) and 'params' in config:
+            # regular hydra-based instantiation
+            instance = hydra.utils.instantiate(config=config)
+        # Hydra 1.x API
+        elif '_target_' in config:
+            # regular hydra-based instantiation
+            instance = hydra.utils.instantiate(config=config)
+        else:
+            instance = None
+            imported_cls_tb = None
+            # Attempt class path resolution from config `target` class (if it exists)
+            if 'target' in config:
+                target_cls = config.target
+                imported_cls = None
+                try:
+                    # try to import the target class
+                    imported_cls = import_class_by_path(target_cls)
+                except Exception:
+                    imported_cls_tb = traceback.format_exc()
+
+                # try instantiating model with target class
+                if imported_cls is not None:
+                    # if calling class (cls) is subclass of imported class,
+                    # use subclass instead
+                    if issubclass(cls, imported_cls):
+                        imported_cls = cls
+
+                    try:
+                        instance = imported_cls(cfg=config)
+                    except Exception:
+                        imported_cls_tb = traceback.format_exc()
+                        instance = None
+
+            # target class resolution was unsuccessful, fall back to current `cls`
+            if instance is None:
+                if imported_cls_tb is not None:
+                    logging.debug(
+                        f"Model instantiation from target class {target_cls} failed with following error.\n"
+                        f"Falling back to `cls`.\n"
+                        f"{imported_cls_tb}"
+                    )
+                instance = cls(cfg=config)
+
+        if not hasattr(instance, '_cfg'):
+            instance._cfg = config
+        return instance
+
+    def to_config_dict(self) -> DictConfig:
+        """Returns object's configuration to config dictionary"""
+        if hasattr(self, '_cfg') and self._cfg is not None and isinstance(self._cfg, DictConfig):
+            # Resolve the config dict
+            config = OmegaConf.to_container(self._cfg, resolve=True)
+            config = OmegaConf.create(config)
+            OmegaConf.set_struct(config, True)
+
+            config = maybe_update_config_version(config)
+
+            self._cfg = config
+
+            return self._cfg
+        else:
+            raise NotImplementedError(
+                'to_config_dict() can currently only return object._cfg but current object does not have it.'
+            )
+
+
+class FileIO(ABC):
+    def save_to(self, save_path: str):
+        """Saves module/model with weights"""
+        raise NotImplementedError()
+
+    @classmethod
+    def restore_from(
+        cls,
+        restore_path: str,
+        override_config_path: Optional[str] = None,
+        map_location: Optional['torch.device'] = None,
+        strict: bool = True,
+        return_config: bool = False,
+    ):
+        """Restores module/model with weights"""
+        raise NotImplementedError()
+
+    @classmethod
+    def from_config_file(cls, path2yaml_file: str):
+        """
+        Instantiates an instance of NeMo Model from YAML config file.
+        Weights will be initialized randomly.
+        Args:
+            path2yaml_file: path to yaml file with model configuration
+
+        Returns:
+
+        """
+        if issubclass(cls, Serialization):
+            conf = OmegaConf.load(path2yaml_file)
+            return cls.from_config_dict(config=conf)
+        else:
+            raise NotImplementedError()
+
+    def to_config_file(self, path2yaml_file: str):
+        """
+        Saves current instance's configuration to YAML config file. Weights will not be saved.
+        Args:
+            path2yaml_file: path2yaml_file: path to yaml file where model model configuration will be saved
+
+        Returns:
+        """
+        if hasattr(self, '_cfg'):
+            self._cfg = maybe_update_config_version(self._cfg)
+            with open(path2yaml_file, 'w') as fout:
+                OmegaConf.save(config=self._cfg, f=fout, resolve=True)
+        else:
+            raise NotImplementedError()
+
+
+@total_ordering
+@dataclass
+class PretrainedModelInfo:
+    pretrained_model_name: str
+    description: str
+    location: str
+    class_: 'Model' = None
+    aliases: List[str] = None
+
+    def __repr__(self):
+        base = self.__class__.__name__
+        extras = (
+            "pretrained_model_name={pretrained_model_name},\n\t"
+            "description={description},\n\t"
+            "location={location}".format(**self.__dict__)
+        )
+
+        if self.class_ is not None:
+            extras = "{extras},\n\t" "class_={class_}".format(extras=extras, **self.__dict__)
+
+        representation = f"{base}(\n\t{extras}\n)"
+        return representation
+
+    def __hash__(self):
+        # assumes that locations are unique urls, and therefore their hashes
+        # should ideally also be unique
+        location_hash = hash(self.location)
+        return location_hash
+
+    def __eq__(self, other):
+        # another object is equal to self, iff
+        # if it's hash is equal to hash(self)
+        return hash(self) == hash(other) or self.pretrained_model_name == other.pretrained_model_name
+
+    def __lt__(self, other):
+        return self.pretrained_model_name < other.pretrained_model_name
+
+
+class Model(Typing, Serialization, FileIO):
+    """
+    Abstract class offering interface which should be implemented by all NeMo models.
+    """
+
+    @classmethod
+    @abstractmethod
+    def list_available_models(cls) -> Optional[PretrainedModelInfo]:
+        """
+        Should list all pre-trained models available via NVIDIA NGC cloud.
+        Note: There is no check that requires model names and aliases to be unique. In the case of a collIsion, whatever
+        model (or alias) is listed first in the this returned list will be instantiated.
+
+        Returns:
+            A list of PretrainedModelInfo entries
+        """
+        pass
+
+    @classmethod
+    def get_available_model_names(cls) -> List[str]:
+        """
+        Returns the list of model names available via NVIDIA NGC cloud,
+        to get the complete model description use list_available_models()
+        Returns:
+            A list of model names
+        """
+        model_names = []
+        if cls.list_available_models() is not None:
+            model_names = [model.pretrained_model_name for model in cls.list_available_models()]
+        return model_names
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name: str,
+        refresh_cache: bool = False,
+        override_config_path: Optional[str] = None,
+        map_location: Optional['torch.device'] = None,
+        strict: bool = True,
+        return_config: bool = False,
+    ):
+        """
+        Instantiates an instance of NeMo from NVIDIA NGC cloud
+        Use restore_from() to instantiate from a local .nemo file.
+        Args:
+            model_name: string key which will be used to find the module.
+            refresh_cache: If set to True, then when fetching from cloud, this will re-fetch the file
+                from cloud even if it is already found in a cache locally.
+            override_config_path: path to a yaml config that will override the internal
+                config file
+            map_location: Optional torch.device() to map the instantiated model to a device.
+                By default (None), it will select a GPU if available, falling back to CPU otherwise.
+            strict: Passed to torch.load_state_dict. By default true.
+            return_config: If set to true, will return just the underlying config of the restored
+                model as an OmegaConf DictConfig object without instantiating the model.
+
+        Returns:
+            A model instance of a particular model class or its underlying config (if return_config is set).
+        """
+        location_in_the_cloud = None
+        description = None
+        models = cls.list_available_models()
+        if models is not None:
+            for pretrained_model_info in cls.list_available_models():
+                found = False
+                if pretrained_model_info.pretrained_model_name == model_name:
+                    found = True
+                elif pretrained_model_info.aliases is not None:
+                    for alias in pretrained_model_info.aliases:
+                        if alias == model_name:
+                            found = True
+                            break
+                if found:
+                    location_in_the_cloud = pretrained_model_info.location
+                    description = pretrained_model_info.description
+                    class_ = pretrained_model_info.class_
+                    break
+
+        if location_in_the_cloud is None:
+            raise FileNotFoundError(
+                f"Model {model_name} was not found. Check cls.list_available_models() for the list of all available models."
+            )
+        filename = location_in_the_cloud.split("/")[-1]
+        url = location_in_the_cloud.replace(filename, "")
+        cache_dir = Path.joinpath(Path.home(), f'.cache/torch/NeMo/NeMo_{nemo.__version__}/{filename[:-5]}')
+        # If either description and location in the cloud changes, this will force re-download
+        cache_subfolder = hashlib.md5((location_in_the_cloud + description).encode('utf-8')).hexdigest()
+        # if file exists on cache_folder/subfolder, it will be re-used, unless refresh_cache is True
+        nemo_model_file_in_cache = maybe_download_from_cloud(
+            url=url, filename=filename, cache_dir=cache_dir, subfolder=cache_subfolder, refresh_cache=refresh_cache
+        )
+        logging.info("Instantiating model from pre-trained checkpoint")
+        if class_ is None:
+            class_ = cls
+        instance = class_.restore_from(
+            restore_path=nemo_model_file_in_cache,
+            override_config_path=override_config_path,
+            map_location=map_location,
+            strict=strict,
+            return_config=return_config,
+        )
+        return instance
+
+
+class typecheck:
+    class TypeState(Enum):
+        """
+        Placeholder to denote the default value of type information provided.
+        If the constructor of this decorator is used to override the class level type definition,
+        this enum value indicate that types will be overridden.
+        """
+
+        UNINITIALIZED = 0
+
+    def __init__(
+        self,
+        input_types: Union[TypeState, Dict[str, NeuralType]] = TypeState.UNINITIALIZED,
+        output_types: Union[TypeState, Dict[str, NeuralType]] = TypeState.UNINITIALIZED,
+        ignore_collections: bool = False,
+    ):
+        """
+        A decorator which performs input-output neural type checks, and attaches
+        neural types to the output of the function that it wraps.
+
+        Requires that the class inherit from `nemo.core.Typing` in order to perform
+        type checking, and will raise an error if that is not the case.
+
+        # Usage (Class level type support)
+        @typecheck()
+        def fn(self, arg1, arg2, ...):
+            ...
+
+        # Usage (Function level type support)
+        @typecheck(input_types=..., output_types=...)
+        def fn(self, arg1, arg2, ...):
+            ...
+
+        Points to be noted:
+        1) The brackets () in `@typecheck()` are necessary.
+
+            You will encounter a TypeError: __init__() takes 1 positional argument but X
+            were given without those brackets.
+
+        2) The function can take any number of positional arguments during definition.
+
+            When you call this function, all arguments must be passed using kwargs only.
+
+        """
+        self.input_types = input_types
+        self.output_types = output_types
+
+        if input_types == self.TypeState.UNINITIALIZED:
+            self.input_override = False
+        else:
+            self.input_override = True
+
+        if output_types == self.TypeState.UNINITIALIZED:
+            self.output_override = False
+        else:
+            self.output_override = True
+
+        self.ignore_collections = ignore_collections
+
+    @wrapt.decorator(enabled=is_typecheck_enabled)
+    def __call__(self, wrapped, instance: Typing, args, kwargs):
+        if instance is None:
+            raise RuntimeError("Only classes which inherit nemo.core.Typing can use this decorator !")
+
+        if not isinstance(instance, Typing):
+            raise RuntimeError("Only classes which inherit nemo.core.Typing can use this decorator !")
+
+        if hasattr(instance, 'input_ports') or hasattr(instance, 'output_ports'):
+            raise RuntimeError(
+                "Typing requires override of `input_types()` and `output_types()`, "
+                "not `input_ports() and `output_ports()`"
+            )
+
+        # Preserve type information
+        if self.input_types is typecheck.TypeState.UNINITIALIZED:
+            self.input_types = instance.input_types
+
+        if self.output_types is typecheck.TypeState.UNINITIALIZED:
+            self.output_types = instance.output_types
+
+        # Resolve global type or local overridden type
+        if self.input_override:
+            input_types = self.input_types
+        else:
+            input_types = instance.input_types
+
+        if self.output_override:
+            output_types = self.output_types
+        else:
+            output_types = instance.output_types
+
+        # If types are not defined, skip type checks and just call the wrapped method
+        if input_types is None and output_types is None:
+            return wrapped(*args, **kwargs)
+
+        # Check that all arguments are kwargs
+        if input_types is not None and len(args) > 0:
+            raise TypeError("All arguments must be passed by kwargs only for typed methods")
+
+        # Perform rudimentary input checks here
+        instance._validate_input_types(input_types=input_types, ignore_collections=self.ignore_collections, **kwargs)
+
+        # Call the method - this can be forward, or any other callable method
+        outputs = wrapped(*args, **kwargs)
+
+        instance._attach_and_validate_output_types(
+            output_types=output_types, ignore_collections=self.ignore_collections, out_objects=outputs
+        )
+
+        return outputs
+
+    @staticmethod
+    def set_typecheck_enabled(enabled: bool = True):
+        global _TYPECHECK_ENABLED
+        _TYPECHECK_ENABLED = enabled
+
+    @staticmethod
+    @contextmanager
+    def disable_checks():
+        typecheck.set_typecheck_enabled(enabled=False)
+        try:
+            yield
+        finally:
+            typecheck.set_typecheck_enabled(enabled=True)

tts/nemo/core/classes/dataset.py

@@ -0,0 +1,109 @@
+# 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
+
+from torch.utils import data
+
+from nemo.core.classes import Serialization, Typing, typecheck
+
+__all__ = ['Dataset', 'IterableDataset']
+
+
+class Dataset(data.Dataset, Typing, Serialization):
+    """Dataset with output ports
+
+    Please Note: Subclasses of IterableDataset should *not* implement input_types.
+    """
+
+    def _collate_fn(self, batch):
+        """
+        A default implementation of a collation function.
+        Users should override this method to define custom data loaders.
+        """
+        return data.dataloader.default_collate(batch)
+
+    @typecheck()
+    def collate_fn(self, batch):
+        """
+        This is the method that user pass as functor to DataLoader.
+        The method optionally performs neural type checking and add types to the outputs.
+
+        Please note, subclasses of Dataset should not implement `input_types`.
+
+        # Usage:
+        dataloader = torch.utils.data.DataLoader(
+                ....,
+                collate_fn=dataset.collate_fn,
+                ....
+        )
+
+        Returns:
+            Collated batch, with or without types.
+        """
+        if self.input_types is not None:
+            raise TypeError("Datasets should not implement `input_types` as they are not checked")
+
+        # Simply forward the inner `_collate_fn`
+        return self._collate_fn(batch)
+
+
+class IterableDataset(data.IterableDataset, Typing, Serialization):
+    """Iterable Dataset with output ports
+
+    Please Note: Subclasses of IterableDataset should *not* implement input_types.
+    """
+
+    def _collate_fn(self, batch):
+        """
+        A default implementation of a collation function.
+        Users should override this method to define custom data loaders.
+        """
+        return data.dataloader.default_collate(batch)
+
+    @typecheck()
+    def collate_fn(self, batch):
+        """
+        This is the method that user pass as functor to DataLoader.
+        The method optionally performs neural type checking and add types to the outputs.
+
+        # Usage:
+        dataloader = torch.utils.data.DataLoader(
+                ....,
+                collate_fn=dataset.collate_fn,
+                ....
+        )
+
+        Returns:
+            Collated batch, with or without types.
+        """
+        if self.input_types is not None:
+            raise TypeError("Datasets should not implement `input_types` as they are not checked")
+
+        # Simply forward the inner `_collate_fn`
+        return self._collate_fn(batch)
+
+
+@dataclass
+class DatasetConfig:
+    """
+
+    """
+
+    # ...
+    batch_size: int = 32
+    drop_last: bool = False
+    shuffle: bool = False
+    num_workers: Optional[int] = None
+    pin_memory: bool = True

tts/nemo/core/classes/exportable.py

@@ -0,0 +1,308 @@
+# 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 os
+from abc import ABC
+from collections import defaultdict
+from enum import Enum
+from typing import Dict
+
+import onnx
+import torch
+
+from nemo.core.classes import typecheck
+from nemo.core.neural_types import AxisKind, NeuralType
+from nemo.utils import logging
+from nemo.utils.export_utils import replace_for_export
+
+__all__ = ['ExportFormat', 'Exportable']
+
+
+class ExportFormat(Enum):
+    """Which format to use when exporting a Neural Module for deployment"""
+
+    ONNX = (1,)
+    TORCHSCRIPT = (2,)
+
+
+_EXT_DICT = {
+    ".pt": ExportFormat.TORCHSCRIPT,
+    ".ts": ExportFormat.TORCHSCRIPT,
+    ".onnx": ExportFormat.ONNX,
+}
+
+
+def get_input_names(self):
+    if not (hasattr(self, 'input_types')):
+        raise NotImplementedError('For export to work you must define input_types')
+    input_names = list(self.input_types.keys())
+    return input_names
+
+
+def get_output_names(self):
+    if not (hasattr(self, 'output_types')):
+        raise NotImplementedError('For export to work you must define output_types')
+    output_names = list(self.output_types.keys())
+    return output_names
+
+
+def get_input_dynamic_axes(self, input_names):
+    dynamic_axes = defaultdict(list)
+    for name in input_names:
+        dynamic_axes = {
+            **dynamic_axes,
+            **Exportable._extract_dynamic_axes(name, self.input_types[name]),
+        }
+    return dynamic_axes
+
+
+def get_output_dynamic_axes(self, output_names):
+    dynamic_axes = defaultdict(list)
+    for name in output_names:
+        dynamic_axes = {
+            **dynamic_axes,
+            **Exportable._extract_dynamic_axes(name, self.output_types[name]),
+        }
+    return dynamic_axes
+
+
+def to_onnxrt_input(input_names, input_list, input_dict):
+    odict = {}
+    for k, v in input_dict.items():
+        odict[k] = v.cpu().numpy()
+    for i, input in enumerate(input_list):
+        odict[input_names[i]] = input.cpu().numpy()
+    return odict
+
+
+class Exportable(ABC):
+    """
+    This Interface should be implemented by particular classes derived from nemo.core.NeuralModule or nemo.core.ModelPT.
+    It gives these entities ability to be exported for deployment to formats such as ONNX.
+    """
+
+    @staticmethod
+    def get_format(filename: str):
+        _, ext = os.path.splitext(filename)
+        try:
+            return _EXT_DICT[ext]
+        except KeyError:
+            raise ValueError(f"Export file {filename} extension does not correspond to any export format!")
+
+    @property
+    def input_module(self):
+        return self
+
+    @property
+    def output_module(self):
+        return self
+
+    def export(
+        self,
+        output: str,
+        input_example=None,
+        output_example=None,
+        verbose=False,
+        export_params=True,
+        do_constant_folding=True,
+        keep_initializers_as_inputs=False,
+        onnx_opset_version: int = 13,
+        try_script: bool = False,
+        set_eval: bool = True,
+        check_trace: bool = False,
+        use_dynamic_axes: bool = True,
+        dynamic_axes=None,
+        check_tolerance=0.01,
+    ):
+        my_args = locals()
+        del my_args['self']
+
+        qual_name = self.__module__ + '.' + self.__class__.__qualname__
+        output_descr = qual_name + ' exported to ONNX'
+
+        try:
+            # Disable typechecks
+            typecheck.set_typecheck_enabled(enabled=False)
+
+            # Allow user to completely override forward method to export
+            if hasattr(type(self), "forward_for_export"):
+                forward_method = type(self).forward_for_export
+                old_forward_method = type(self).forward
+                type(self).forward = forward_method
+            else:
+                forward_method = None
+
+            # Set module to eval mode
+            if set_eval:
+                self.freeze()
+                self.input_module.freeze()
+                self.output_module.freeze()
+
+            format = self.get_format(output)
+
+            if input_example is None:
+                input_example = self.input_module.input_example()
+
+            my_args['input_example'] = input_example
+
+            # Run (posibly overridden) prepare method before calling forward()
+            self._prepare_for_export(**my_args)
+
+            input_list = list(input_example)
+            input_dict = {}
+            # process possible kwargs
+            if isinstance(input_list[-1], dict):
+                input_dict = input_list[-1]
+                input_list = input_list[:-1]
+
+            input_names = get_input_names(self.input_module)
+            # remove unnecessary inputs for input_ports
+            for name in self.disabled_deployment_input_names:
+                input_names.remove(name)
+            output_names = get_output_names(self.output_module)
+            # remove unnecessary inputs for input_ports
+            for name in self.disabled_deployment_output_names:
+                output_names.remove(name)
+
+            output_example = self.forward(*input_list, **input_dict)
+
+            with torch.jit.optimized_execution(True), torch.no_grad():
+                jitted_model = None
+                if try_script:
+                    try:
+                        jitted_model = torch.jit.script(self)
+                    except Exception as e:
+                        print("jit.script() failed!", e)
+
+                if format == ExportFormat.TORCHSCRIPT:
+                    if jitted_model is None:
+                        jitted_model = torch.jit.trace_module(
+                            self,
+                            {"forward": tuple(input_list) + tuple(input_dict.values())},
+                            strict=False,
+                            optimize=True,
+                            check_trace=check_trace,
+                            check_tolerance=check_tolerance,
+                        )
+                    if verbose:
+                        print(jitted_model.code)
+                    jitted_model.save(output)
+                    assert os.path.exists(output)
+
+                elif format == ExportFormat.ONNX:
+                    if jitted_model is None:
+                        jitted_model = self
+
+                    # dynamic axis is a mapping from input/output_name => list of "dynamic" indices
+                    if dynamic_axes is None and use_dynamic_axes:
+                        dynamic_axes = get_input_dynamic_axes(self.input_module, input_names)
+                        dynamic_axes = {**dynamic_axes, **get_output_dynamic_axes(self.output_module, output_names)}
+
+                    torch.onnx.export(
+                        jitted_model,
+                        input_example,
+                        output,
+                        input_names=input_names,
+                        output_names=output_names,
+                        verbose=verbose,
+                        export_params=export_params,
+                        do_constant_folding=do_constant_folding,
+                        keep_initializers_as_inputs=keep_initializers_as_inputs,
+                        dynamic_axes=dynamic_axes,
+                        opset_version=onnx_opset_version,
+                        example_outputs=output_example,
+                    )
+
+                    # Verify the model can be read, and is valid
+                    onnx_model = onnx.load(output)
+                    onnx.checker.check_model(onnx_model, full_check=True)
+                    test_runtime = check_trace
+                    if test_runtime:
+                        try:
+                            import onnxruntime
+                        except (ImportError, ModuleNotFoundError):
+                            test_runtime = False
+                            logging.warning(
+                                f"ONNX generated at {output}, not verified - please install onnxruntime.\n"
+                            )
+
+                    if test_runtime:
+                        sess = onnxruntime.InferenceSession(onnx_model.SerializeToString())
+
+                        ort_out = sess.run(None, to_onnxrt_input(input_names, input_list, input_dict))
+                        all_good = True
+                        for out_name, out in enumerate(ort_out):
+                            expected = output_example[out_name].cpu()
+                            if not torch.allclose(
+                                torch.from_numpy(out), expected, rtol=check_tolerance, atol=100 * check_tolerance
+                            ):
+                                all_good = False
+                                logging.info(
+                                    f"onnxruntime results mismatch! PyTorch(expected):\n{expected}\nONNXruntime:\n{out}"
+                                )
+                        status = "SUCCESS" if all_good else "FAIL"
+                        logging.info(f"ONNX generated at {output} verified with onnxruntime : " + status)
+                else:
+                    raise ValueError(f'Encountered unknown export format {format}.')
+        finally:
+            typecheck.set_typecheck_enabled(enabled=True)
+            if forward_method:
+                type(self).forward = old_forward_method
+        return ([output], [output_descr])
+
+    @property
+    def disabled_deployment_input_names(self):
+        """Implement this method to return a set of input names disabled for export"""
+        return set()
+
+    @property
+    def disabled_deployment_output_names(self):
+        """Implement this method to return a set of output names disabled for export"""
+        return set()
+
+    @property
+    def supported_export_formats(self):
+        """Implement this method to return a set of export formats supported. Default is all types."""
+        return set([ExportFormat.ONNX, ExportFormat.TORCHSCRIPT])
+
+    @staticmethod
+    def _extract_dynamic_axes(name: str, ntype: NeuralType):
+        """
+        Implement this method to provide dynamic axes id for ONNX export.
+        By default, this method will extract BATCH and TIME dimension ids from each provided input/output name argument.
+
+        For example, if module/model accepts argument named "input_signal" with type corresponding to [Batch, Time, Dim]
+        shape, then the returned result should contain "input_signal" -> [0, 1] because Batch and Time are dynamic axes
+        as they can change from call to call during inference.
+
+        Args:
+            name: Name of input or output parameter
+            ntype: Corresponding Neural Type
+
+        Returns:
+
+        """
+        dynamic_axes = defaultdict(list)
+        if ntype.axes:
+            for ind, axis in enumerate(ntype.axes):
+                if axis.kind in [AxisKind.Batch, AxisKind.Time, AxisKind.Width, AxisKind.Height]:
+                    dynamic_axes[name].append(ind)
+        return dynamic_axes
+
+    def _prepare_for_export(self, **kwargs):
+        """
+        Override this method to prepare module for export. This is in-place operation.
+        Base version does common necessary module replacements (Apex etc)
+        """
+        replace_1D_2D = kwargs.get('replace_1D_2D', False)
+        replace_for_export(self, replace_1D_2D)

tts/nemo/core/classes/loss.py

@@ -0,0 +1,26 @@
+# 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 torch
+
+from nemo.core.classes.common import Serialization, Typing
+
+__all__ = ['Loss']
+
+
+class Loss(torch.nn.modules.loss._Loss, Typing, Serialization):
+    """Inherit this class to implement custom loss."""
+
+    def __init__(self, **kwargs):
+        super(Loss, self).__init__(**kwargs)

tts/nemo/core/classes/module.py

@@ -0,0 +1,77 @@
+# 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 contextlib import contextmanager
+
+from torch.nn import Module
+
+from nemo.core.classes.common import FileIO, Serialization, Typing
+
+__all__ = ['NeuralModule']
+
+
+class NeuralModule(Module, Typing, Serialization, FileIO):
+    """
+    Abstract class offering interface shared between all PyTorch Neural Modules.
+    """
+
+    @property
+    def num_weights(self):
+        """
+        Utility property that returns the total number of parameters of NeuralModule.
+        """
+        num: int = 0
+        for p in self.parameters():
+            if p.requires_grad:
+                num += p.numel()
+        return num
+
+    def input_example(self):
+        """
+        Override this method if random inputs won't work
+        Returns:
+            A tuple sample of valid input data.
+        """
+
+        return
+
+    def freeze(self) -> None:
+        r"""
+        Freeze all params for inference.
+        """
+        for param in self.parameters():
+            param.requires_grad = False
+
+        self.eval()
+
+    def unfreeze(self) -> None:
+        """
+        Unfreeze all parameters for training.
+        """
+        for param in self.parameters():
+            param.requires_grad = True
+
+        self.train()
+
+    @contextmanager
+    def as_frozen(self):
+        """
+        Context manager which temporarily freezes a module, yields control and finally unfreezes the module.
+        """
+        self.freeze()
+
+        try:
+            yield
+        finally:
+            self.unfreeze()

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,
+}

tts/nemo/core/neural_types/__init__.py

@@ -0,0 +1,19 @@
+# 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.neural_types.axes import *
+from nemo.core.neural_types.comparison import *
+from nemo.core.neural_types.elements import *
+from nemo.core.neural_types.neural_type import *

tts/nemo/core/neural_types/axes.py

@@ -0,0 +1,101 @@
+# 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 enum import Enum
+from typing import Optional
+
+__all__ = ['AxisKindAbstract', 'AxisKind', 'AxisType']
+
+
+class AxisKindAbstract(Enum):
+    """This is an abstract Enum to represents what does varying axis dimension mean.
+    In practice, you will almost always use AxisKind Enum. This Enum should be inherited by
+    your OWN Enum if you aren't satisfied with AxisKind. Then your own Enum can be used
+    instead of AxisKind."""
+
+    pass
+
+
+class AxisKind(AxisKindAbstract):
+    """This Enum represents what does varying axis dimension mean.
+    For example, does this dimension correspond to width, batch, time, etc.
+    The "Dimension" and "Channel" kinds are the same and used to represent
+    a general axis. "Any" axis will accept any axis kind fed to it.
+    """
+
+    Batch = 0
+    Time = 1
+    Dimension = 2
+    Channel = 2
+    Width = 3
+    Height = 4
+    Any = 5
+    Sequence = 6
+    FlowGroup = 7
+    Singleton = 8  # Used to represent a axis that has size 1
+
+    def __repr__(self):
+        return self.__str__()
+
+    def __str__(self):
+        return str(self.name).lower()
+
+    @staticmethod
+    def from_str(label):
+        """Returns AxisKind instance based on short string representation"""
+        _label = label.lower().strip()
+        if _label == "b" or _label == "n" or _label == "batch":
+            return AxisKind.Batch
+        elif _label == "t" or _label == "time":
+            return AxisKind.Time
+        elif _label == "d" or _label == "c" or _label == "channel":
+            return AxisKind.Dimension
+        elif _label == "w" or _label == "width":
+            return AxisKind.Width
+        elif _label == "h" or _label == "height":
+            return AxisKind.Height
+        elif _label == "s" or _label == "singleton":
+            return AxisKind.Singleton
+        elif _label == "flowgroup":
+            return AxisKind.FlowGroup
+        elif _label == "any":
+            return AxisKind.Any
+        else:
+            raise ValueError(f"Can't create AxisKind from {label}")
+
+
+class AxisType(object):
+    """This class represents axis semantics and (optionally) it's dimensionality
+       Args:
+           kind (AxisKindAbstract): what kind of axis it is? For example Batch, Height, etc.
+           size (int, optional): specify if the axis should have a fixed size. By default it is set to None and you
+           typically do not want to set it for Batch and Time
+           is_list (bool, default=False): whether this is a list or a tensor axis
+    """
+
+    def __init__(self, kind: AxisKindAbstract, size: Optional[int] = None, is_list=False):
+        if size is not None and is_list:
+            raise ValueError("The axis can't be list and have a fixed size")
+        self.kind = kind
+        self.size = size
+        self.is_list = is_list
+
+    def __repr__(self):
+        if self.size is None:
+            representation = str(self.kind)
+        else:
+            representation = f"{str(self.kind)}:{self.size}"
+        if self.is_list:
+            representation += "_listdim"
+        return representation

tts/nemo/core/neural_types/comparison.py

@@ -0,0 +1,32 @@
+# 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 enum import Enum
+
+__all__ = ['NeuralTypeComparisonResult']
+
+
+class NeuralTypeComparisonResult(Enum):
+    """The result of comparing two neural type objects for compatibility.
+    When comparing A.compare_to(B):"""
+
+    SAME = 0
+    LESS = 1  # A is B
+    GREATER = 2  # B is A
+    DIM_INCOMPATIBLE = 3  # Resize connector might fix incompatibility
+    TRANSPOSE_SAME = 4  # A transpose and/or converting between lists and tensors will make them same
+    CONTAINER_SIZE_MISMATCH = 5  # A and B contain different number of elements
+    INCOMPATIBLE = 6  # A and B are incompatible
+    SAME_TYPE_INCOMPATIBLE_PARAMS = 7  # A and B are of the same type but parametrized differently
+    UNCHECKED = 8  # type comparison wasn't done

tts/nemo/core/neural_types/elements.py

@@ -0,0 +1,324 @@
+# 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 abc
+from abc import ABC
+from typing import Dict, Optional, Tuple
+
+from nemo.core.neural_types.comparison import NeuralTypeComparisonResult
+
+__all__ = [
+    'ElementType',
+    'VoidType',
+    'ChannelType',
+    'AcousticEncodedRepresentation',
+    'AudioSignal',
+    'SpectrogramType',
+    'MelSpectrogramType',
+    'MFCCSpectrogramType',
+    'LogitsType',
+    'LabelsType',
+    'HypothesisType',
+    'LossType',
+    'RegressionValuesType',
+    'CategoricalValuesType',
+    'PredictionsType',
+    'LogprobsType',
+    'LengthsType',
+    'EmbeddedTextType',
+    'EncodedRepresentation',
+    'MaskType',
+    'Target',
+    'ClassificationTarget',
+    'ImageFeatureValue',
+    'Index',
+    'ImageValue',
+    'NormalizedImageValue',
+    'StringLabel',
+    'StringType',
+    'TokenIndex',
+    'Length',
+    'IntType',
+    'FloatType',
+    'NormalDistributionSamplesType',
+    'NormalDistributionMeanType',
+    'NormalDistributionLogVarianceType',
+    'TokenDurationType',
+    'TokenLogDurationType',
+    'LogDeterminantType',
+    'SequenceToSequenceAlignmentType',
+]
+
+
+class ElementType(ABC):
+    """Abstract class defining semantics of the tensor elements.
+    We are relying on Python for inheritance checking"""
+
+    def __str__(self):
+        return self.__doc__
+
+    def __repr__(self):
+        return self.__class__.__name__
+
+    @property
+    def type_parameters(self) -> Dict:
+        """Override this property to parametrize your type. For example, you can specify 'storage' type such as
+        float, int, bool with 'dtype' keyword. Another example, is if you want to represent a signal with a
+        particular property (say, sample frequency), then you can put sample_freq->value in there.
+        When two types are compared their type_parameters must match."""
+        return {}
+
+    @property
+    def fields(self) -> Optional[Tuple]:
+        """This should be used to logically represent tuples/structures. For example, if you want to represent a
+        bounding box (x, y, width, height) you can put a tuple with names ('x', y', 'w', 'h') in here.
+        Under the hood this should be converted to the last tesnor dimension of fixed size = len(fields).
+        When two types are compared their fields must match."""
+        return None
+
+    def compare(self, second) -> NeuralTypeComparisonResult:
+        # First, check general compatibility
+        first_t = type(self)
+        second_t = type(second)
+
+        if first_t == second_t:
+            result = NeuralTypeComparisonResult.SAME
+        elif issubclass(first_t, second_t):
+            result = NeuralTypeComparisonResult.LESS
+        elif issubclass(second_t, first_t):
+            result = NeuralTypeComparisonResult.GREATER
+        else:
+            result = NeuralTypeComparisonResult.INCOMPATIBLE
+
+        if result != NeuralTypeComparisonResult.SAME:
+            return result
+        else:
+            # now check that all parameters match
+            check_params = set(self.type_parameters.keys()) == set(second.type_parameters.keys())
+            if check_params is False:
+                return NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+            else:
+                for k1, v1 in self.type_parameters.items():
+                    if v1 is None or second.type_parameters[k1] is None:
+                        # Treat None as Void
+                        continue
+                    if v1 != second.type_parameters[k1]:
+                        return NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+            # check that all fields match
+            if self.fields == second.fields:
+                return NeuralTypeComparisonResult.SAME
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
+
+
+class VoidType(ElementType):
+    """Void-like type which is compatible with everything.
+    It is a good practice to use this type only as necessary.
+    For example, when you need template-like functionality.
+    """
+
+    def compare(cls, second: abc.ABCMeta) -> NeuralTypeComparisonResult:
+        return NeuralTypeComparisonResult.SAME
+
+
+# TODO: Consider moving these files elsewhere
+class ChannelType(ElementType):
+    """Element to represent convolutional input/output channel.
+    """
+
+
+class EmbeddedTextType(ChannelType):
+    """Element to represent output on word/text embedding layers
+    """
+
+
+class LogitsType(ElementType):
+    """Element type to represent logits"""
+
+
+class ProbsType(ElementType):
+    """Element type to represent probabilities. For example, outputs of softmax layers."""
+
+
+class LogprobsType(ElementType):
+    """Element type to represent log-probabilities. For example, outputs of log softmax layers."""
+
+
+class LabelsType(ElementType):
+    """Element type to represent some sort of labels. This is often used as a base class to create
+    a more concrete types such as RegressionValuesType, etc."""
+
+
+class HypothesisType(LabelsType):
+    """Element type to represent some decoded hypothesis, which may further be processed to obtain
+    a concrete label."""
+
+
+class LengthsType(ElementType):
+    """Element type representing lengths of something"""
+
+
+class LossType(ElementType):
+    """Element type to represent outputs of Loss modules"""
+
+
+class EncodedRepresentation(ChannelType):
+    """Element type to represent encoded representation, for example, encoder's output"""
+
+
+class AcousticEncodedRepresentation(EncodedRepresentation):
+    """Element type to represent encoded representation returned by the acoustic encoder model"""
+
+
+class AudioSignal(ElementType):
+    """Element type to represent encoded representation returned by the acoustic encoder model
+    Args:
+        freq (int): sampling frequency of a signal. Note that two signals will only be the same if their
+        freq is the same.
+    """
+
+    def __init__(self, freq: int = None):
+        self._params = {}
+        self._params['freq'] = freq
+
+    @property
+    def type_parameters(self):
+        return self._params
+
+
+class SpectrogramType(ChannelType):
+    """Element type to represent generic spectrogram signal"""
+
+
+class MelSpectrogramType(SpectrogramType):
+    """Element type to represent mel spectrogram signal"""
+
+
+class MFCCSpectrogramType(SpectrogramType):
+    """Element type to represent MFCC spectrogram signal"""
+
+
+class PredictionsType(LabelsType):
+    """Element type to represent some sort of predictions returned by model"""
+
+
+class RegressionValuesType(PredictionsType):
+    """Element type to represent labels for regression task"""
+
+
+class CategoricalValuesType(PredictionsType):
+    """Element type to represent labels for categorical classification task"""
+
+
+class MaskType(PredictionsType):
+    """Element type to represent a boolean mask"""
+
+
+class Index(ElementType):
+    """Type representing an element being an index of the sample."""
+
+
+class Target(ElementType):
+    """
+        Type representing an element being a target value.
+    """
+
+
+class ClassificationTarget(Target):
+    """
+        Type representing an element being target value in the classification task, i.e. identifier of a desired class.
+    """
+
+
+class ImageValue(ElementType):
+    """
+        Type representing an element/value of a single image channel,
+        e.g. a single element (R) of RGB image.
+    """
+
+
+class NormalizedImageValue(ImageValue):
+    """
+        Type representing an element/value of a single image channel normalized to <0-1> range,
+        e.g. a single element (R) of normalized RGB image.
+    """
+
+
+class ImageFeatureValue(ImageValue):
+    """Type representing an element (single value) of a (image) feature maps."""
+
+
+class StringType(ElementType):
+    """Element type representing a single string"""
+
+
+class StringLabel(StringType):
+    """
+        Type representing an label being a string with class name (e.g. the "hamster" class in CIFAR100).
+    """
+
+
+class BoolType(ElementType):
+    """Element type representing a single integer"""
+
+
+class IntType(ElementType):
+    """Element type representing a single integer"""
+
+
+class FloatType(ElementType):
+    """Element type representing a single float"""
+
+
+class TokenIndex(IntType):
+    """Type representing an element being index of a token in some kind of a vocabulary."""
+
+
+class Length(IntType):
+    """Type representing an element storing a "length" (e.g. length of a list)."""
+
+
+class ProbabilityDistributionSamplesType(ElementType):
+    """Element to represent tensors that meant to be sampled from a valid probability distribution
+    """
+
+
+class NormalDistributionSamplesType(ProbabilityDistributionSamplesType):
+    """Element to represent tensors that meant to be sampled from a valid normal distribution
+    """
+
+
+class SequenceToSequenceAlignmentType(ElementType):
+    """Class to represent the alignment from seq-to-seq attention outputs. Generally a mapping from endcoder time steps
+    to decoder time steps."""
+
+
+class NormalDistributionMeanType(ElementType):
+    """Element to represent the mean of a normal distribution"""
+
+
+class NormalDistributionLogVarianceType(ElementType):
+    """Element to represent the log variance of a normal distribution"""
+
+
+class TokenDurationType(ElementType):
+    """Element for representing the duration of a token"""
+
+
+class TokenLogDurationType(ElementType):
+    """Element for representing the log-duration of a token"""
+
+
+class LogDeterminantType(ElementType):
+    """Element for representing log determinants usually used in flow models"""

tts/nemo/core/neural_types/neural_type.py

@@ -0,0 +1,243 @@
+# 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 typing import Optional, Tuple
+
+from nemo.core.neural_types.axes import AxisKind, AxisType
+from nemo.core.neural_types.comparison import NeuralTypeComparisonResult
+from nemo.core.neural_types.elements import ElementType, VoidType
+
+__all__ = [
+    'NeuralType',
+    'NeuralTypeError',
+    'NeuralPortNameMismatchError',
+    'NeuralPortNmTensorMismatchError',
+]
+
+
+class NeuralType(object):
+    """This is the main class which would represent neural type concept.
+    It is used to represent *the types* of inputs and outputs.
+
+    Args:
+        axes (Optional[Tuple]): a tuple of AxisTypes objects representing the semantics of what varying each axis means
+            You can use a short, string-based form here. For example: ('B', 'C', 'H', 'W') would correspond to an NCHW
+            format frequently used in computer vision. ('B', 'T', 'D') is frequently used for signal processing and
+            means [batch, time, dimension/channel].
+        elements_type (ElementType): an instance of ElementType class representing the semantics of what is stored
+            inside the tensor. For example: logits (LogitsType), log probabilities (LogprobType), etc.
+        optional (bool): By default, this is false. If set to True, it would means that input to the port of this
+            type can be optional.
+    """
+
+    def __str__(self):
+
+        if self.axes is not None:
+            return f"axes: {self.axes}; elements_type: {self.elements_type.__class__.__name__}"
+        else:
+            return f"axes: None; elements_type: {self.elements_type.__class__.__name__}"
+
+    def __init__(self, axes: Optional[Tuple] = None, elements_type: ElementType = VoidType(), optional=False):
+        if not isinstance(elements_type, ElementType):
+            raise ValueError(
+                "elements_type of NeuralType must be an instance of a class derived from ElementType. "
+                "Did you pass a class instead?"
+            )
+        self.elements_type = elements_type
+        if axes is not None:
+            NeuralType.__check_sanity(axes)
+            axes_list = []
+            for axis in axes:
+                if isinstance(axis, str):
+                    axes_list.append(AxisType(AxisKind.from_str(axis), None))
+                elif isinstance(axis, AxisType):
+                    axes_list.append(axis)
+                else:
+                    raise ValueError("axis type must be either str or AxisType instance")
+            self.axes = tuple(axes_list)
+        else:
+            self.axes = None
+        self.optional = optional
+
+    def compare(self, second) -> NeuralTypeComparisonResult:
+        """Performs neural type comparison of self with second. When you chain two modules' inputs/outputs via
+        __call__ method, this comparison will be called to ensure neural type compatibility."""
+        # First, handle dimensionality
+        axes_a = self.axes
+        axes_b = second.axes
+
+        # "Big void" type
+        if isinstance(self.elements_type, VoidType) and self.axes is None:
+            return NeuralTypeComparisonResult.SAME
+
+        if self.axes is None:
+            if second.axes is None:
+                return self.elements_type.compare(second.elements_type)
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
+
+        dimensions_pass = NeuralType.__compare_axes(axes_a, axes_b)
+        element_comparison_result = self.elements_type.compare(second.elements_type)
+
+        # SAME DIMS
+        if dimensions_pass == 0:
+            return element_comparison_result
+        # TRANSPOSE_SAME DIMS
+        elif dimensions_pass == 1:
+            if element_comparison_result == NeuralTypeComparisonResult.SAME:
+                return NeuralTypeComparisonResult.TRANSPOSE_SAME
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
+        # DIM_INCOMPATIBLE DIMS
+        elif dimensions_pass == 2:
+            if element_comparison_result == NeuralTypeComparisonResult.SAME:
+                return NeuralTypeComparisonResult.DIM_INCOMPATIBLE
+            else:
+                return NeuralTypeComparisonResult.INCOMPATIBLE
+        else:
+            return NeuralTypeComparisonResult.INCOMPATIBLE
+
+    def compare_and_raise_error(self, parent_type_name, port_name, second_object):
+        """ Method compares definition of one type with another and raises an error if not compatible. """
+        type_comatibility = self.compare(second_object)
+        if (
+            type_comatibility != NeuralTypeComparisonResult.SAME
+            and type_comatibility != NeuralTypeComparisonResult.GREATER
+        ):
+            raise NeuralPortNmTensorMismatchError(
+                parent_type_name, port_name, str(self), str(second_object.ntype), type_comatibility
+            )
+
+    def __eq__(self, other):
+        if isinstance(other, NeuralType):
+            return self.compare(other)
+
+        return False
+
+    @staticmethod
+    def __check_sanity(axes):
+        # check that list come before any tensor dimension
+        are_strings = True
+        for axis in axes:
+            if not isinstance(axis, str):
+                are_strings = False
+            if isinstance(axis, str) and not are_strings:
+                raise ValueError("Either use full class names or all strings")
+        if are_strings:
+            return
+        checks_passed = True
+        saw_tensor_dim = False
+        for axis in axes:
+            if not axis.is_list:
+                saw_tensor_dim = True
+            else:  # current axis is a list
+                if saw_tensor_dim:  # which is preceded by tensor dim
+                    checks_passed = False
+        if not checks_passed:
+            raise ValueError(
+                "You have list dimension after Tensor dimension. All list dimensions must preceed Tensor dimensions"
+            )
+
+    @staticmethod
+    def __compare_axes(axes_a, axes_b) -> int:
+        """
+        Compares axes_a and axes_b
+        Args:
+            axes_a: first axes tuple
+            axes_b: second axes tuple
+
+        Returns:
+            0 - if they are exactly the same
+            1 - if they are "TRANSPOSE_SAME"
+            2 - if the are "DIM_INCOMPATIBLE"
+            3 - if they are different
+        """
+        if axes_a is None and axes_b is None:
+            return 0
+        elif axes_a is None and axes_b is not None:
+            return 3
+        elif axes_a is not None and axes_b is None:
+            return 3
+        elif len(axes_a) != len(axes_b):
+            return 3
+        # After these ifs we know that len(axes_a) == len(axes_b)
+
+        same = True
+        kinds_a = dict()
+        kinds_b = dict()
+        for axis_a, axis_b in zip(axes_a, axes_b):
+            kinds_a[axis_a.kind] = axis_a.size
+            kinds_b[axis_b.kind] = axis_b.size
+            if axis_a.kind == AxisKind.Any:
+                same = True
+            elif (
+                axis_a.kind != axis_b.kind
+                or axis_a.is_list != axis_b.is_list
+                or (axis_a.size != axis_b.size and axis_a.size is not None)
+            ):
+                same = False
+        if same:
+            return 0
+        else:
+            # can be TRANSPOSE_SAME, DIM_INCOMPATIBLE
+            if kinds_a.keys() == kinds_b.keys():
+                for key, value in kinds_a.items():
+                    if kinds_b[key] != value:
+                        return 2
+                return 1
+            else:
+                return 3
+
+    def __repr__(self):
+        if self.axes is not None:
+            axes = str(self.axes)
+        else:
+            axes = "None"
+
+        if self.elements_type is not None:
+            element_type = repr(self.elements_type)
+        else:
+            element_type = "None"
+
+        data = f"axis={axes}, element_type={element_type}"
+
+        if self.optional:
+            data = f"{data}, optional={self.optional}"
+
+        final = f"{self.__class__.__name__}({data})"
+        return final
+
+
+class NeuralTypeError(Exception):
+    """Base class for neural type related exceptions."""
+
+
+class NeuralPortNameMismatchError(NeuralTypeError):
+    """Exception raised when neural module is called with incorrect port
+    names."""
+
+    def __init__(self, input_port_name):
+        super().__init__()
+        self.message = "Wrong input port name: {0}".format(input_port_name)
+
+
+class NeuralPortNmTensorMismatchError(NeuralTypeError):
+    """Exception raised when a port is fed with a NmTensor of incompatible
+    type."""
+
+    def __init__(self, class_name, port_name, first_type, second_type, type_comatibility):
+        super().__init__()
+        self.message = "\nIn {}. \nPort: {} and a NmTensor it was fed are \n".format(class_name, port_name)
+        self.message += "of incompatible neural types:\n\n{} \n\n and \n\n{}".format(first_type, second_type)
+        self.message += "\n\nType comparison result: {}".format(type_comatibility)

tts/nemo/core/optim/__init__.py

@@ -0,0 +1,29 @@
+# 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.optim.lr_scheduler import (
+    CosineAnnealing,
+    InverseSquareRootAnnealing,
+    NoamAnnealing,
+    PolynomialDecayAnnealing,
+    PolynomialHoldDecayAnnealing,
+    SquareAnnealing,
+    SquareRootAnnealing,
+    WarmupAnnealing,
+    WarmupHoldPolicy,
+    WarmupPolicy,
+    prepare_lr_scheduler,
+)
+from nemo.core.optim.novograd import Novograd
+from nemo.core.optim.optimizers import get_optimizer, parse_optimizer_args, register_optimizer

tts/nemo/core/optim/lr_scheduler.py

@@ -0,0 +1,686 @@
+# 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 copy
+import dataclasses
+import math
+import warnings
+from functools import partial
+from typing import Any, Dict, Optional, Union
+
+import hydra
+import torch.optim as optim
+import torch.optim.lr_scheduler as pt_scheduler
+import torch.utils.data.dataloader as dataloader
+from omegaconf import DictConfig, OmegaConf
+from torch.optim.lr_scheduler import _LRScheduler
+
+from nemo.core.config import SchedulerParams, get_scheduler_config, register_scheduler_params
+from nemo.utils import logging
+from nemo.utils.model_utils import maybe_update_config_version
+
+
+class WarmupPolicy(_LRScheduler):
+    """Adds warmup kwargs and warmup logic to lr policy.
+    All arguments should be passed as kwargs for clarity,
+    Args:
+        warmup_steps: Number of training steps in warmup stage
+        warmup_ratio: Ratio of warmup steps to total steps
+        max_steps: Total number of steps while training or `None` for
+            infinite training
+    """
+
+    def __init__(self, optimizer, *, warmup_steps=None, warmup_ratio=None, max_steps=None, min_lr=0.0, last_epoch=-1):
+        assert not (
+            warmup_steps is not None and warmup_ratio is not None
+        ), "Either use particular number of step or ratio"
+        assert warmup_ratio is None or max_steps is not None, "If there is a ratio, there should be a total steps"
+
+        # It is necessary to assign all attributes *before* __init__,
+        # as class is wrapped by an inner class.
+        self.max_steps = max_steps
+        if warmup_steps is not None:
+            self.warmup_steps = warmup_steps
+        elif warmup_ratio is not None:
+            self.warmup_steps = int(warmup_ratio * max_steps)
+        else:
+            self.warmup_steps = 0
+
+        self.min_lr = min_lr
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if not self._get_lr_called_within_step:
+            warnings.warn(
+                "To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
+            )
+
+        step = self.last_epoch
+
+        if step <= self.warmup_steps:
+            lr_val = (step + 1) / (self.warmup_steps + 1)
+            return [initial_lr * lr_val for initial_lr in self.base_lrs]
+
+        if step > self.max_steps:
+            return [self.min_lr for _ in self.base_lrs]
+
+        return self._get_lr(step)
+
+    def _get_lr(self, step):
+        """Simple const lr policy"""
+        return self.base_lrs
+
+
+class WarmupHoldPolicy(WarmupPolicy):
+    """Variant of WarmupPolicy which maintains high learning rate for a defined number of steps.
+    All arguments should be passed as kwargs for clarity,
+    Args:
+        warmup_steps: Number of training steps in warmup stage
+        warmup_ratio: Ratio of warmup steps to total steps
+        hold_steps: Number of training steps to hold the learning rate after warm up
+        hold_ratio: Ratio of hold steps to total steps
+        max_steps: Total number of steps while training or `None` for
+            infinite training
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        *,
+        warmup_steps=None,
+        warmup_ratio=None,
+        hold_steps=None,
+        hold_ratio=None,
+        max_steps=None,
+        min_lr=0.0,
+        last_epoch=-1,
+    ):
+        assert not (hold_steps is not None and hold_ratio is not None), "Either use particular number of step or ratio"
+        assert hold_ratio is None or max_steps is not None, "If there is a ratio, there should be a total steps"
+
+        self.min_lr = min_lr
+        self._last_warmup_lr = 0.0
+
+        # Necessary to duplicate as class attributes are hidden in inner class
+        self.max_steps = max_steps
+        if warmup_steps is not None:
+            self.warmup_steps = warmup_steps
+        elif warmup_ratio is not None:
+            self.warmup_steps = int(warmup_ratio * max_steps)
+        else:
+            self.warmup_steps = 0
+
+        if hold_steps is not None:
+            self.hold_steps = hold_steps + self.warmup_steps
+        elif hold_ratio is not None:
+            self.hold_steps = int(hold_ratio * max_steps) + self.warmup_steps
+        else:
+            self.hold_steps = 0
+
+        super().__init__(
+            optimizer,
+            warmup_steps=warmup_steps,
+            warmup_ratio=warmup_ratio,
+            max_steps=max_steps,
+            last_epoch=last_epoch,
+            min_lr=min_lr,
+        )
+
+    def get_lr(self):
+        if not self._get_lr_called_within_step:
+            warnings.warn(
+                "To get the last learning rate computed by the scheduler, " "please use `get_last_lr()`.", UserWarning
+            )
+
+        step = self.last_epoch
+
+        # Warmup phase
+        if step <= self.warmup_steps:
+            lr_val = (step + 1) / (self.warmup_steps + 1)
+            return [initial_lr * lr_val for initial_lr in self.base_lrs]
+
+        # Hold phase
+        if (step >= self.warmup_steps) and (step < self.hold_steps):
+            return self.base_lrs
+
+        if step > self.max_steps:
+            return [self.min_lr for _ in self.base_lrs]
+
+        return self._get_lr(step)
+
+
+def _squareroot_annealing(initial_lr, step, max_steps, min_lr):
+    mult = ((max_steps - step) / max_steps) ** 0.5
+    out_lr = initial_lr * mult
+    out_lr = max(out_lr, min_lr)
+    return out_lr
+
+
+def _square_annealing(initial_lr, step, max_steps, min_lr):
+    mult = ((max_steps - step) / max_steps) ** 2
+    out_lr = initial_lr * mult
+    out_lr = max(out_lr, min_lr)
+    return out_lr
+
+
+def _cosine_annealing(initial_lr, step, max_steps, min_lr):
+    mult = 0.5 * (1 + math.cos(math.pi * step / max_steps))
+    out_lr = (initial_lr - min_lr) * mult + min_lr
+    return out_lr
+
+
+def _poly_decay(initial_lr, step, decay_steps, power, min_lr, cycle):
+    if cycle:
+        multiplier = 1.0 if step == 0 else math.ceil(step / decay_steps)
+        decay_steps *= multiplier
+    else:
+        step = min(step, decay_steps)
+    p = step / decay_steps
+    lr = (initial_lr - min_lr) * math.pow(1.0 - p, power)
+    lr += min_lr
+    return lr
+
+
+class SquareAnnealing(WarmupPolicy):
+    def __init__(self, optimizer, *, max_steps, min_lr=1e-5, last_epoch=-1, **kwargs):
+        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
+
+    def _get_lr(self, step):
+        new_lrs = [
+            _square_annealing(
+                initial_lr=initial_lr,
+                step=step - self.warmup_steps,
+                max_steps=self.max_steps - self.warmup_steps,
+                min_lr=self.min_lr,
+            )
+            for initial_lr in self.base_lrs
+        ]
+        return new_lrs
+
+
+class SquareRootAnnealing(WarmupPolicy):
+    def __init__(self, optimizer, *, max_steps, min_lr=0, last_epoch=-1, **kwargs):
+        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
+
+    def _get_lr(self, step):
+        new_lrs = [
+            _squareroot_annealing(initial_lr=initial_lr, step=step, max_steps=self.max_steps, min_lr=self.min_lr)
+            for initial_lr in self.base_lrs
+        ]
+        return new_lrs
+
+
+class CosineAnnealing(WarmupPolicy):
+    def __init__(self, optimizer, *, max_steps, min_lr=0, last_epoch=-1, **kwargs):
+        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
+
+    def _get_lr(self, step):
+        for initial_lr in self.base_lrs:
+            if initial_lr < self.min_lr:
+                raise ValueError(
+                    f"{self} received an initial learning rate that was lower than the minimum learning rate."
+                )
+
+        new_lrs = [
+            _cosine_annealing(
+                initial_lr=initial_lr,
+                step=step - self.warmup_steps,
+                max_steps=self.max_steps - self.warmup_steps,
+                min_lr=self.min_lr,
+            )
+            for initial_lr in self.base_lrs
+        ]
+        return new_lrs
+
+
+class NoamAnnealing(_LRScheduler):
+    def __init__(
+        self, optimizer, *, d_model, warmup_steps=None, warmup_ratio=None, max_steps=None, min_lr=0.0, last_epoch=-1
+    ):
+        self._normalize = d_model ** (-0.5)
+        assert not (
+            warmup_steps is not None and warmup_ratio is not None
+        ), "Either use particular number of step or ratio"
+        assert warmup_ratio is None or max_steps is not None, "If there is a ratio, there should be a total steps"
+
+        # It is necessary to assign all attributes *before* __init__,
+        # as class is wrapped by an inner class.
+        self.max_steps = max_steps
+        if warmup_steps is not None:
+            self.warmup_steps = warmup_steps
+        elif warmup_ratio is not None:
+            self.warmup_steps = int(warmup_ratio * max_steps)
+        else:
+            self.warmup_steps = 0
+
+        self.min_lr = min_lr
+        super().__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if not self._get_lr_called_within_step:
+            warnings.warn(
+                "To get the last learning rate computed by the scheduler, please use `get_last_lr()`.", UserWarning
+            )
+
+        step = max(1, self.last_epoch)
+
+        if step > self.max_steps:
+            return [self.min_lr for _ in self.base_lrs]
+
+        for initial_lr in self.base_lrs:
+            if initial_lr < self.min_lr:
+                raise ValueError(
+                    f"{self} received an initial learning rate that was lower than the minimum learning rate."
+                )
+
+        new_lrs = [self._noam_annealing(initial_lr=initial_lr, step=step) for initial_lr in self.base_lrs]
+        return new_lrs
+
+    def _noam_annealing(self, initial_lr, step):
+        mult = self._normalize * min(step ** (-0.5), step * (self.warmup_steps ** (-1.5)))
+        out_lr = initial_lr * mult
+        if step > self.warmup_steps:
+            out_lr = max(out_lr, self.min_lr)
+        return out_lr
+
+
+class WarmupAnnealing(WarmupPolicy):
+    def __init__(self, optimizer, *, max_steps, last_epoch=-1, min_lr=0.0, **kwargs):
+        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
+
+    def _get_lr(self, step):
+        progress = float(step / self.max_steps)
+        warmup_ratio = float(self.warmup_steps / self.max_steps)
+
+        mult = max((progress - 1.0) / (warmup_ratio - 1.0), 0.0)
+        out_lr = [initial_lr * mult for initial_lr in self.base_lrs]
+
+        return out_lr
+
+
+class InverseSquareRootAnnealing(WarmupPolicy):
+    def __init__(self, optimizer, *, max_steps, last_epoch=-1, min_lr=0.0, **kwargs):
+        super().__init__(optimizer=optimizer, max_steps=max_steps, **kwargs, last_epoch=last_epoch, min_lr=min_lr)
+
+    def _get_lr(self, step):
+        denom = ((step + 1) / (self.warmup_steps + 1)) ** 0.5
+        out_lr = [initial_lr / denom for initial_lr in self.base_lrs]
+        return out_lr
+
+
+class PolynomialDecayAnnealing(WarmupPolicy):
+    def __init__(self, optimizer, *, max_steps, min_lr=0.0, power=1.0, cycle=False, last_epoch=-1, **kwargs):
+        self.power = power
+        self.cycle = cycle
+
+        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
+
+    def _get_lr(self, step):
+        new_lrs = [
+            _poly_decay(
+                initial_lr,
+                step=step - self.warmup_steps,
+                decay_steps=self.max_steps - self.warmup_steps,
+                power=self.power,
+                min_lr=self.min_lr,
+                cycle=self.cycle,
+            )
+            for initial_lr in self.base_lrs
+        ]
+        return new_lrs
+
+
+class PolynomialHoldDecayAnnealing(WarmupHoldPolicy):
+    def __init__(self, optimizer, *, max_steps, min_lr=0.0, power=1.0, cycle=False, last_epoch=-1, **kwargs):
+        self.power = power
+        self.cycle = cycle
+
+        super().__init__(optimizer=optimizer, max_steps=max_steps, last_epoch=last_epoch, min_lr=min_lr, **kwargs)
+
+    def _get_lr(self, step):
+        new_lrs = [
+            _poly_decay(
+                initial_lr,
+                step=step - self.hold_steps,
+                decay_steps=self.max_steps - max(self.warmup_steps, self.hold_steps),
+                power=self.power,
+                min_lr=self.min_lr,
+                cycle=self.cycle,
+            )
+            for initial_lr in self.base_lrs
+        ]
+        return new_lrs
+
+
+def register_scheduler(name: str, scheduler: _LRScheduler, scheduler_params: SchedulerParams):
+    """
+    Checks if the scheduler 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: Scheduler class (inherits from _LRScheduler)
+        scheduler_params: The parameters as a dataclass of the scheduler
+    """
+    if name in AVAILABLE_SCHEDULERS:
+        raise ValueError(f"Cannot override pre-existing schedulers. Conflicting scheduler name = {name}")
+
+    AVAILABLE_SCHEDULERS[name] = scheduler
+
+    sched_name = "{}_params".format(scheduler.__name__)
+    register_scheduler_params(name=sched_name, scheduler_params=scheduler_params)
+
+
+def get_scheduler(name: str, **kwargs: Optional[Dict[str, Any]]) -> _LRScheduler:
+    """
+    Convenience method to obtain an _LRScheduler class and partially instantiate it with optimizer kwargs.
+
+    Args:
+        name: Name of the scheduler in the registry.
+        kwargs: Optional kwargs of the scheduler used during instantiation.
+
+    Returns:
+        a partially instantiated _LRScheduler
+    """
+    if name not in AVAILABLE_SCHEDULERS:
+        raise ValueError(
+            f"Cannot resolve scheduler{name}'. Available optimizers are : " f"{AVAILABLE_SCHEDULERS.keys()}"
+        )
+
+    scheduler_cls = AVAILABLE_SCHEDULERS[name]
+    scheduler = partial(scheduler_cls, **kwargs)
+    return scheduler
+
+
+def prepare_lr_scheduler(
+    optimizer: optim.Optimizer,
+    scheduler_config: Union[Dict[str, Any], DictConfig],
+    train_dataloader: Optional[dataloader.DataLoader] = None,
+) -> Optional[Dict[str, Any]]:
+    """
+    Constructs an LR Scheduler (optionally) for a given optimizer, based on a config with the following schema
+
+    optim:
+      name: <name of optimizer>
+      lr: <maximal learning rate>
+
+      # <additional optimizer arguments>
+      args:
+        name: auto  # special keyword, resolves to correct optimizer config for given optimizer name
+        # cls: nemo.core.config.optimizers.NovogradParams  # explicit instantiation by class path
+        params:  # optional override parameters for the optimizer config
+          betas: [0.8, 0.5]
+          weight_decay: 0.001
+
+      # scheduler setup
+      sched:
+        name: <name of scheduler>
+        iters_per_batch: null # computed at runtime; mandatory to have
+        max_steps: null # computed at runtime or explicitly set here; mandatory to have
+
+        # pytorch lightning args <mandatory>
+        monitor: val_loss
+        reduce_on_plateau: false
+
+        # <scheduler config override>
+        args:
+          name: auto  # special keyword, resolves to correct optimizer config for given optimizer name
+          # cls: nemo.core.config.schedulers.CosineAnnealingParams  # explicit instantiation by class path
+          params:  # optional override parameters for the optimizer config
+            warmup_steps: null
+            warmup_ratio: null
+            min_lr: 0.0
+            last_epoch: -1
+
+    Args:
+        optimizer: An instantiated Optimizer.
+        scheduler_config: A dictionary / config dict which follows the above schema.
+        train_dataloader: Optional requirement, must be passed if "iters_per_batch" is defined
+            instead of "max_steps". Used to compute effective "max_steps".
+
+    Returns:
+        A dictionary containing the LR Scheduler implementation if the config was successfully parsed
+        along with other parameters required by Pytorch Lightning, otherwise None.
+    """
+    if scheduler_config is not None:
+        scheduler_config = maybe_update_config_version(scheduler_config)
+
+    # Build nested dictionary for convenience out of structured objects
+    if isinstance(scheduler_config, DictConfig):
+        scheduler_config = OmegaConf.to_container(scheduler_config, resolve=True)
+
+    elif dataclasses.is_dataclass(scheduler_config):
+        # Recursively transform data classes to basic dictionaries
+        scheduler_config = OmegaConf.create(scheduler_config)
+        scheduler_config = OmegaConf.to_container(scheduler_config, resolve=True)
+
+    # Test to see if config follows above schema
+
+    add_max_args_flag = True
+    interval = 'step'
+    if scheduler_config is not None:
+        if 'args' in scheduler_config:
+            scheduler_args = scheduler_config.pop('args')
+        else:
+            scheduler_args = copy.deepcopy(scheduler_config)
+
+            # Remove extra parameters from scheduler_args nest
+            # Assume all other parameters are to be passed into scheduler constructor
+
+            if 'name' in scheduler_args and scheduler_args['name'] == 'ReduceLROnPlateau':
+                add_max_args_flag = False
+                interval = 'epoch'
+
+            scheduler_args.pop('name', None)
+            scheduler_args.pop('t_max_epochs', None)
+            scheduler_args.pop('t_accumulate_grad_batches', None)
+            scheduler_args.pop('t_limit_train_batches', None)
+            scheduler_args.pop('t_num_workers', None)
+            scheduler_args.pop('monitor', None)
+            scheduler_args.pop('reduce_on_plateau', None)
+
+    else:
+        # Return gracefully in case `sched` was not supplied; inform user
+        logging.info('Scheduler not initialized as no `sched` config supplied to setup_optimizer()')
+        return None
+
+    # Try instantiation of scheduler params from config class path
+    if '_target_' in scheduler_args:
+        scheduler_args_cfg = OmegaConf.create(scheduler_args)
+        scheduler_conf = hydra.utils.instantiate(scheduler_args_cfg)
+        scheduler_args = vars(scheduler_conf)
+
+        # Get name of the scheduler
+        scheduler_name = scheduler_conf.__class__.__name__
+
+        if 'Params' in scheduler_name:
+            scheduler_name = scheduler_name.replace('Params', '')
+
+    else:
+        # Class path instantiation failed; try resolving "name" component
+
+        # Get name of the scheduler
+        if 'name' in scheduler_config:
+            scheduler_name = scheduler_config['name']
+        else:
+            logging.warning(
+                "Could not resolve classpath for Scheduler Config, and `name` "
+                "was not provided either. \n"
+                "Scheduler cannot be instantiated !"
+            )
+            return None
+
+        # If class path was not provided, perhaps `name` is provided for resolution
+        if 'name' in scheduler_args:
+            # If `auto` is passed as name for resolution of optimizer name,
+            # then lookup optimizer name and resolve its parameter config
+            if scheduler_args['name'] == 'auto':
+                scheduler_params_name = "{}Params".format(scheduler_name)
+            else:
+                scheduler_params_name = scheduler_args['name']
+
+            # Get override arguments provided in the config yaml file / Dict Config
+            scheduler_params_override = scheduler_args.get('params', {})
+
+            # If params is itself a dict config object provided explicitly in Dict Config
+            # Resolve to dictionary for convenience
+            if isinstance(scheduler_params_override, DictConfig):
+                scheduler_params_override = OmegaConf.to_container(scheduler_params_override, resolve=True)
+
+            # Get and instantiate the Config dataclass for this scheduler
+            scheduler_params_cls = get_scheduler_config(scheduler_params_name, **scheduler_params_override)
+            scheduler_params = scheduler_params_cls()  # instantiate the parameters object
+            scheduler_args = vars(scheduler_params)  # extract just the dictionary from the Config object
+
+        else:
+            # assume the input dictionary is schedular args (from dataclasses / omegaconf)
+            pass
+
+    # Extract value to monitor in losses, if provided.
+    if 'monitor' in scheduler_config:
+        monitor = scheduler_config.get('monitor')
+    else:
+        # Default to train loss
+        monitor = 'loss'
+
+    # Store exact max_steps if it is provided
+    if 'max_steps' in scheduler_config and scheduler_config['max_steps'] is not None:
+        max_steps = scheduler_config['max_steps']
+
+    elif 't_max_epochs' in scheduler_config:
+        # Compute effective max_steps if t_max_epochs is provided
+        if train_dataloader is None:
+            logging.warning(
+                'As `t_max_epochs` is provided/computed, it is required to pass the train dataloader in order\n'
+                'to compute effective maximum number of steps.\n'
+                'Scheduler will not be instantiated !'
+            )
+            return None
+
+        # Raise exception if neither `max_steps` nor `t_max_epochs` is provided
+        if scheduler_config.get('t_max_epochs', None) is None:
+            logging.warning(
+                "`t_max_epochs` cannot be None when `max_steps` is not not provided.\n"
+                "This can occur when `train dataloader` is not available to correctly "
+                "prepare the scheduler.\n"
+                "Scheduler will not be instantiated !"
+            )
+            return None
+
+        # Get iters_per_batch
+        max_epochs = scheduler_config.get('t_max_epochs')
+        accumulate_grad_batches = scheduler_config.get('t_accumulate_grad_batches')
+        limit_train_batches = scheduler_config.get('t_limit_train_batches')
+        num_workers = scheduler_config.get('t_num_workers')
+
+        # Compute effective num max_steps
+        num_samples = len(train_dataloader.dataset)
+        batch_size = train_dataloader.batch_size
+        drop_last = train_dataloader.drop_last
+
+        max_steps = compute_max_steps(
+            max_epochs=max_epochs,
+            accumulate_grad_batches=accumulate_grad_batches,
+            limit_train_batches=limit_train_batches,
+            num_workers=num_workers,
+            num_samples=num_samples,
+            batch_size=batch_size,
+            drop_last=drop_last,
+        )
+
+    else:
+        logging.warning(
+            "Neither `max_steps` nor `iters_per_batch` were provided to `optim.sched`, "
+            "cannot compute effective `max_steps` !\n"
+            "Scheduler will not be instantiated !"
+        )
+        return None
+
+    # Inject max_steps (effective or provided) into the scheduler config
+    if add_max_args_flag:
+        scheduler_args['max_steps'] = max_steps
+
+    # Get the scheduler class from the config
+    scheduler_cls = get_scheduler(scheduler_name, **scheduler_args)
+
+    # Instantiate the LR schedule
+    schedule = scheduler_cls(optimizer, **scheduler_args)
+
+    logging.info(
+        'Scheduler "%s" \nwill be used during training (effective maximum steps = %d) - \nParameters : \n(%s)',
+        str(schedule),
+        max_steps,
+        OmegaConf.to_yaml(OmegaConf.create(scheduler_args)),
+    )
+
+    # Wrap the schedule in PTL arguments to perform stepwise computation
+    # Rather than epoch level computation
+    if isinstance(schedule, optim.lr_scheduler.ReduceLROnPlateau):
+        reduce_lr_on_plateau = True
+    else:
+        reduce_lr_on_plateau = False
+
+    schedule_dict = {
+        'scheduler': schedule,
+        'interval': interval,
+        'frequency': 1,
+        'monitor': monitor,
+        'reduce_on_plateau': reduce_lr_on_plateau,
+    }
+    return schedule_dict
+
+
+def compute_max_steps(
+    max_epochs, accumulate_grad_batches, limit_train_batches, num_workers, num_samples, batch_size, drop_last
+):
+    _round = math.floor if drop_last else math.ceil
+
+    sampler_num_samples = math.ceil(num_samples / num_workers)
+
+    if drop_last and num_workers > 1:
+        logging.warning(
+            "Please note that drop_last is broken in pytorch 1.6.0. We will fix when pytorch 1.7.0 is released"
+        )
+        # TODO: Master verion, not in pytorch 1.6.0
+        # sampler_num_samples = math.ceil((num_samples - num_workers)/ num_workers)
+
+    steps_per_epoch = _round(sampler_num_samples / batch_size)
+    if isinstance(limit_train_batches, int) or limit_train_batches == 0.0:
+        steps_per_epoch = min(steps_per_epoch, int(limit_train_batches))
+    elif steps_per_epoch != float('inf'):
+        # limit_train_batches is a percentage of batches per epoch
+        steps_per_epoch = int(steps_per_epoch * limit_train_batches)
+        if accumulate_grad_batches == 1:
+            steps_per_epoch = max(steps_per_epoch, 1)
+
+    return math.ceil(steps_per_epoch / accumulate_grad_batches) * max_epochs
+
+
+AVAILABLE_SCHEDULERS = {
+    'WarmupPolicy': WarmupPolicy,
+    'WarmupHoldPolicy': WarmupHoldPolicy,
+    'SquareAnnealing': SquareAnnealing,
+    'CosineAnnealing': CosineAnnealing,
+    'NoamAnnealing': NoamAnnealing,
+    'WarmupAnnealing': WarmupAnnealing,
+    'InverseSquareRootAnnealing': InverseSquareRootAnnealing,
+    'SquareRootAnnealing': SquareRootAnnealing,
+    'PolynomialDecayAnnealing': PolynomialDecayAnnealing,
+    'PolynomialHoldDecayAnnealing': PolynomialHoldDecayAnnealing,
+    'StepLR': pt_scheduler.StepLR,
+    'ExponentialLR': pt_scheduler.ExponentialLR,
+    'ReduceLROnPlateau': pt_scheduler.ReduceLROnPlateau,
+    'CyclicLR': pt_scheduler.CyclicLR,
+}

tts/nemo/core/optim/novograd.py

@@ -0,0 +1,145 @@
+# 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 torch
+from torch.optim.optimizer import Optimizer
+
+__all__ = ['Novograd']
+
+
+def _check_valid_opt_params(lr, eps, betas):
+    if lr < 0:
+        raise ValueError(f"Invalid learning rate: {lr}")
+    if eps < 0:
+        raise ValueError(f"Invalid epsilon value: {eps}")
+    if not (0.0 <= betas[0] < 1.0 and 0.0 <= betas[1] < 1.0):
+        raise ValueError(f"Betas have to be between 0 and 1: {betas}")
+
+
+class Novograd(Optimizer):
+    """Implements Novograd algorithm.
+    It has been proposed  in "Stochastic Gradient Methods with Layer-wise
+    Adaptive Moments for Training of Deep Networks"
+    (https://arxiv.org/abs/1905.11286)
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        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"
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.95, 0.98),
+        eps=1e-8,
+        weight_decay=0,
+        grad_averaging=False,
+        amsgrad=False,
+        luc=False,
+        luc_trust=1e-3,
+        luc_eps=1e-8,
+    ):
+        _check_valid_opt_params(lr, eps, betas)
+        defaults = dict(
+            lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, grad_averaging=grad_averaging, amsgrad=amsgrad,
+        )
+        self.luc = luc
+        self.luc_trust = luc_trust
+        self.luc_eps = luc_eps
+        super(Novograd, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Novograd, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("amsgrad", False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("Sparse gradients are not supported.")
+                amsgrad = group["amsgrad"]
+                state = self.state[p]
+
+                # State initialization
+                if not state:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros([]).to(state["exp_avg"].device)
+                    if amsgrad:
+                        # Maintains max of all exp moving avg of squared grad
+                        state["max_exp_avg_sq"] = torch.zeros([]).to(state["exp_avg"].device)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                if amsgrad:
+                    max_exp_avg_sq = state["max_exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                norm = grad.norm().pow(2)
+
+                if exp_avg_sq == 0:
+                    exp_avg_sq.copy_(norm)
+                else:
+                    exp_avg_sq.mul_(beta2).add_(norm, alpha=1.0 - beta2)
+
+                if amsgrad:
+                    # Maintains max of all 2nd moment running avg till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group["eps"])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+
+                grad.div_(denom)
+                if group["weight_decay"] != 0:
+                    grad.add_(p.data, alpha=group["weight_decay"])
+                if group["grad_averaging"]:
+                    grad.mul_(1 - beta1)
+                exp_avg.mul_(beta1).add_(grad)
+
+                if self.luc:
+                    # Clip update so that updates are less than eta*weights
+                    data_norm = torch.norm(p.data)
+                    grad_norm = torch.norm(exp_avg.data)
+                    luc_factor = self.luc_trust * data_norm / (grad_norm + self.luc_eps)
+                    luc_factor = min(luc_factor, group["lr"])
+                    p.data.add_(exp_avg, alpha=-luc_factor)
+                else:
+                    p.data.add_(exp_avg, alpha=-group["lr"])
+
+        return loss

tts/nemo/core/optim/optimizers.py

@@ -0,0 +1,171 @@
+# 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 copy
+from functools import partial
+from typing import Any, Dict, List, Optional, Union
+
+import hydra
+import torch.optim as optim
+from omegaconf import DictConfig, OmegaConf
+from torch.optim import adadelta, adagrad, adamax, rmsprop, rprop
+from torch.optim.optimizer import Optimizer
+
+from nemo.core.config import OptimizerParams, get_optimizer_config, register_optimizer_params
+from nemo.core.optim.novograd import Novograd
+from nemo.utils import logging
+from nemo.utils.model_utils import maybe_update_config_version
+
+AVAILABLE_OPTIMIZERS = {
+    'sgd': optim.SGD,
+    'adam': optim.Adam,
+    'adamw': optim.AdamW,
+    'adadelta': adadelta.Adadelta,
+    'adamax': adamax.Adamax,
+    'adagrad': adagrad.Adagrad,
+    'rmsprop': rmsprop.RMSprop,
+    'rprop': rprop.Rprop,
+    'novograd': Novograd,
+}
+
+try:
+    from apex.optimizers import FusedLAMB
+
+    AVAILABLE_OPTIMIZERS['lamb'] = FusedLAMB
+except ModuleNotFoundError:
+    logging.warning("Apex was not found. Using the lamb optimizer will error out.")
+
+__all__ = ['get_optimizer', 'register_optimizer', 'parse_optimizer_args']
+
+
+def parse_optimizer_args(
+    optimizer_name: str, optimizer_kwargs: Union[DictConfig, Dict[str, Any]]
+) -> Union[Dict[str, Any], DictConfig]:
+    """
+    Parses a list of strings, of the format "key=value" or "key2=val1,val2,..."
+    into a dictionary of type {key=value, key2=[val1, val2], ...}
+
+    This dictionary is then used to instantiate the chosen Optimizer.
+
+    Args:
+        optimizer_name: string name of the optimizer, used for auto resolution of params
+        optimizer_kwargs: Either a list of strings in a specified format,
+            or a dictionary. If a dictionary is provided, it is assumed the dictionary
+            is the final parsed value, and simply returned.
+            If a list of strings is provided, each item in the list is parsed into a
+            new dictionary.
+
+    Returns:
+        A dictionary
+    """
+    kwargs = {}
+
+    if optimizer_kwargs is None:
+        return kwargs
+
+    optimizer_kwargs = copy.deepcopy(optimizer_kwargs)
+    optimizer_kwargs = maybe_update_config_version(optimizer_kwargs)
+
+    if isinstance(optimizer_kwargs, DictConfig):
+        optimizer_kwargs = OmegaConf.to_container(optimizer_kwargs, resolve=True)
+
+    # If it is a dictionary, perform stepwise resolution
+    if hasattr(optimizer_kwargs, 'keys'):
+        # Attempt class path resolution
+        if '_target_' in optimizer_kwargs:  # captures (target, _target_)
+            optimizer_kwargs_config = OmegaConf.create(optimizer_kwargs)
+            optimizer_instance = hydra.utils.instantiate(optimizer_kwargs_config)  # type: DictConfig
+            optimizer_instance = vars(optimizer_instance)
+            return optimizer_instance
+
+        # If class path was not provided, perhaps `name` is provided for resolution
+        if 'name' in optimizer_kwargs:
+            # If `auto` is passed as name for resolution of optimizer name,
+            # then lookup optimizer name and resolve its parameter config
+            if optimizer_kwargs['name'] == 'auto':
+                optimizer_params_name = "{}_params".format(optimizer_name)
+                optimizer_kwargs.pop('name')
+            else:
+                optimizer_params_name = optimizer_kwargs.pop('name')
+
+            # Override arguments provided in the config yaml file
+            if 'params' in optimizer_kwargs:
+                # If optimizer kwarg overrides are wrapped in yaml `params`
+                optimizer_params_override = optimizer_kwargs.get('params')
+            else:
+                # If the kwargs themselves are a DictConfig
+                optimizer_params_override = optimizer_kwargs
+
+            if isinstance(optimizer_params_override, DictConfig):
+                optimizer_params_override = OmegaConf.to_container(optimizer_params_override, resolve=True)
+
+            optimizer_params_cls = get_optimizer_config(optimizer_params_name, **optimizer_params_override)
+
+            # If we are provided just a Config object, simply return the dictionary of that object
+            if optimizer_params_name is None:
+                optimizer_params = vars(optimizer_params_cls)
+                return optimizer_params
+
+            else:
+                # If we are provided a partial class instantiation of a Config,
+                # Instantiate it and retrieve its vars as a dictionary
+                optimizer_params = optimizer_params_cls()  # instantiate the parameters object
+                optimizer_params = vars(optimizer_params)
+                return optimizer_params
+
+        # simply return the dictionary that was provided
+        return optimizer_kwargs
+
+    return kwargs
+
+
+def register_optimizer(name: str, optimizer: Optimizer, optimizer_params: OptimizerParams):
+    """
+    Checks if the optimizer name exists in the registry, and if it doesnt, adds it.
+
+    This allows custom optimizers 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: Optimizer class
+        optimizer_params: The parameters as a dataclass of the optimizer
+    """
+    if name in AVAILABLE_OPTIMIZERS:
+        raise ValueError(f"Cannot override pre-existing optimizers. Conflicting optimizer name = {name}")
+
+    AVAILABLE_OPTIMIZERS[name] = optimizer
+
+    optim_name = "{}_params".format(optimizer.__name__)
+    register_optimizer_params(name=optim_name, optimizer_params=optimizer_params)
+
+
+def get_optimizer(name: str, **kwargs: Optional[Dict[str, Any]]) -> Optimizer:
+    """
+    Convenience method to obtain an Optimizer class and partially instantiate it with optimizer kwargs.
+
+    Args:
+        name: Name of the Optimizer in the registry.
+        kwargs: Optional kwargs of the optimizer used during instantiation.
+
+    Returns:
+        a partially instantiated Optimizer
+    """
+    if name not in AVAILABLE_OPTIMIZERS:
+        raise ValueError(
+            f"Cannot resolve optimizer '{name}'. Available optimizers are : " f"{AVAILABLE_OPTIMIZERS.keys()}"
+        )
+
+    optimizer = AVAILABLE_OPTIMIZERS[name]
+    optimizer = partial(optimizer, **kwargs)
+    return optimizer

tts/nemo/core/utils/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2021, 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.utils import numba_utils

tts/nemo/core/utils/numba_utils.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2021, 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 contextlib
+import logging as pylogger
+import operator
+import os
+
+from nemo.utils import model_utils
+
+# Prevent Numba CUDA logs from showing at info level
+cuda_logger = pylogger.getLogger('numba.cuda.cudadrv.driver')
+cuda_logger.setLevel(pylogger.ERROR)  # only show error
+
+__NUMBA_MINIMUM_VERSION__ = "0.53.0"
+
+
+NUMBA_INSTALLATION_MESSAGE = (
+    "Could not import `numba`.\n"
+    "Please install numba in one of the following ways."
+    "1) If using conda, simply install it with conda using `conda install -c numba numba`\n"
+    "2) If using pip (not recommended), `pip install --upgrade numba`\n"
+    "followed by `export NUMBAPRO_LIBDEVICE='/usr/local/cuda/nvvm/libdevice/'` and \n"
+    "`export NUMBAPRO_NVVM='/usr/local/cuda/nvvm/lib64/libnvvm.so'`.\n"
+    "It is advised to always install numba using conda only, "
+    "as pip installations might interfere with other libraries such as llvmlite.\n"
+    "If pip install does not work, you can also try adding `--ignore-installed` to the pip command,\n"
+    "but this is not advised."
+)
+
+STRICT_NUMBA_COMPAT_CHECK = True
+
+# Get environment key if available
+if 'STRICT_NUMBA_COMPAT_CHECK' in os.environ:
+    check_str = os.environ.get('STRICT_NUMBA_COMPAT_CHECK')
+    check_bool = str(check_str).lower() in ("yes", "true", "t", "1")
+    STRICT_NUMBA_COMPAT_CHECK = check_bool
+
+
+def is_numba_compat_strict() -> bool:
+    """
+    Returns strictness level of numba cuda compatibility checks.
+
+    If value is true, numba cuda compatibility matrix must be satisfied.
+    If value is false, only cuda availability is checked, not compatibility.
+    Numba Cuda may still compile and run without issues in such a case, or it may fail.
+    """
+    return STRICT_NUMBA_COMPAT_CHECK
+
+
+def set_numba_compat_strictness(strict: bool):
+    """
+    Sets the strictness level of numba cuda compatibility checks.
+
+    If value is true, numba cuda compatibility matrix must be satisfied.
+    If value is false, only cuda availability is checked, not compatibility.
+    Numba Cuda may still compile and run without issues in such a case, or it may fail.
+
+    Args:
+        strict: bool value, whether to enforce strict compatibility checks or relax them.
+    """
+    global STRICT_NUMBA_COMPAT_CHECK
+    STRICT_NUMBA_COMPAT_CHECK = strict
+
+
+@contextlib.contextmanager
+def with_numba_compat_strictness(strict: bool):
+    initial_strictness = is_numba_compat_strict()
+    set_numba_compat_strictness(strict=strict)
+    yield
+    set_numba_compat_strictness(strict=initial_strictness)
+
+
+def numba_cuda_is_supported(min_version: str) -> bool:
+    """
+    Tests if an appropriate version of numba is installed, and if it is,
+    if cuda is supported properly within it.
+    
+    Args:
+        min_version: The minimum version of numba that is required.
+
+    Returns:
+        bool, whether cuda is supported with this current installation or not.
+    """
+    module_available, msg = model_utils.check_lib_version('numba', checked_version=min_version, operator=operator.ge)
+
+    # If numba is not installed
+    if module_available is None:
+        return False
+
+    # If numba version is installed and available
+    if module_available is True:
+        from numba import cuda
+
+        # this method first arrived in 0.53, and that's the minimum version required
+        if hasattr(cuda, 'is_supported_version'):
+            try:
+                cuda_available = cuda.is_available()
+                if cuda_available:
+                    cuda_compatible = cuda.is_supported_version()
+                else:
+                    cuda_compatible = False
+
+                if is_numba_compat_strict():
+                    return cuda_available and cuda_compatible
+                else:
+                    return cuda_available
+
+            except OSError:
+                # dlopen(libcudart.dylib) might fail if CUDA was never installed in the first place.
+                return False
+        else:
+            # assume cuda is supported, but it may fail due to CUDA incompatibility
+            return False
+
+    else:
+        return False
+
+
+def skip_numba_cuda_test_if_unsupported(min_version: str):
+    """
+    Helper method to skip pytest test case if numba cuda is not supported.
+    
+    Args:
+        min_version: The minimum version of numba that is required.
+    """
+    numba_cuda_support = numba_cuda_is_supported(min_version)
+    if not numba_cuda_support:
+        import pytest
+
+        pytest.skip(f"Numba cuda test is being skipped. Minimum version required : {min_version}")