Introduce code from Rust monorepo. This is a mess

tts/NeMo/tests/core/test_config_utils.py

@@ -0,0 +1,125 @@
+# 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
+
+import pytest
+import pytorch_lightning as ptl
+
+from nemo.core.config.pytorch_lightning import TrainerConfig
+from nemo.utils import config_utils
+
+
+@pytest.fixture()
+def cls():
+    class DummyClass:
+        def __init__(self, a, b=5, c: int = 0, d: 'ABC' = None):
+            pass
+
+    return DummyClass
+
+
+class TestConfigUtils:
+    @pytest.mark.unit
+    def test_all_args_exist(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+            d: Any = None
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_all_args_dont_exist(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert not signatures_match
+        assert len(cls_subset) > 0
+        assert len(dataclass_subset) == 0
+
+    @pytest.mark.unit
+    def test_extra_args_exist(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+            d: Any = None
+            e: float = 0.0
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert not signatures_match
+        assert len(cls_subset) == 0
+        assert len(dataclass_subset) > 0
+
+    @pytest.mark.unit
+    def test_extra_args_exist_but_is_ignored(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+            d: Any = None
+            e: float = 0.0  # Assume ignored
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass, ignore_args=['e'])
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_args_exist_but_is_remapped(self, cls):
+        @dataclass
+        class DummyDataClass:
+            a: int = -1
+            b: int = 5
+            c: int = 0
+            e: Any = None  # Assume remapped
+
+        result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass, remap_args={'e': 'd'})
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None
+
+    @pytest.mark.unit
+    def test_ptl_config(self):
+        PTL_DEPRECATED = ['distributed_backend', 'automatic_optimization']
+
+        result = config_utils.assert_dataclass_signature_match(ptl.Trainer, TrainerConfig, ignore_args=PTL_DEPRECATED)
+        signatures_match, cls_subset, dataclass_subset = result
+
+        assert signatures_match
+        assert cls_subset is None
+        assert dataclass_subset is None

tts/NeMo/tests/core/test_exp_manager.py

@@ -0,0 +1,357 @@
+# 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 math
+import re
+from pathlib import Path
+
+import pytest
+import pytorch_lightning as pl
+import torch
+from omegaconf import OmegaConf
+from omegaconf.errors import OmegaConfBaseException
+
+from nemo.constants import NEMO_ENV_VARNAME_VERSION
+from nemo.core.classes import ModelPT
+from nemo.utils.exp_manager import (
+    CheckpointMisconfigurationError,
+    LoggerMisconfigurationError,
+    NotFoundError,
+    exp_manager,
+)
+
+
+class MyTestOptimizer(torch.optim.Optimizer):
+    def __init__(self, params):
+        self._step = 0
+        super().__init__(params, {})
+
+    def step(self, *args, **kwargs):
+        for group in self.param_groups:
+            for p in group['params']:
+                if self._step == 0:
+                    p.data = 0.1 * torch.ones(p.shape)
+                elif self._step == 1:
+                    p.data = 0.0 * torch.ones(p.shape)
+                else:
+                    p.data = 0.01 * torch.ones(p.shape)
+        self._step += 1
+        return None
+
+
+class OnesDataset(torch.utils.data.Dataset):
+    def __init__(self, dataset_len):
+        super().__init__()
+        self.__dataset_len = dataset_len
+
+    def __getitem__(self, *args):
+        return torch.ones(2)
+
+    def __len__(self):
+        return self.__dataset_len
+
+
+class ExampleModel(ModelPT):
+    def __init__(self, *args, **kwargs):
+        cfg = OmegaConf.structured({})
+        super().__init__(cfg)
+        self.l1 = torch.nn.modules.Linear(in_features=2, out_features=1)
+
+    def train_dataloader(self):
+        dataset = OnesDataset(2)
+        return torch.utils.data.DataLoader(dataset, batch_size=2)
+
+    def val_dataloader(self):
+        dataset = OnesDataset(10)
+        return torch.utils.data.DataLoader(dataset, batch_size=2)
+
+    def forward(self, batch):
+        output = self.l1(batch)
+        output = torch.nn.functional.l1_loss(output, torch.zeros(output.size()).to(output.device))
+        return output
+
+    def validation_step(self, batch, batch_idx):
+        return self(batch)
+
+    def training_step(self, batch, batch_idx):
+        return self(batch)
+
+    def configure_optimizers(self):
+        return MyTestOptimizer(self.parameters())
+        # return torch.optim.Adam(self.parameters(), lr=0.1)
+
+    def list_available_models(self):
+        pass
+
+    def setup_training_data(self):
+        pass
+
+    def setup_validation_data(self):
+        pass
+
+    def validation_epoch_end(self, loss):
+        self.log("val_loss", torch.stack(loss).mean())
+
+
+class TestExpManager:
+    @pytest.mark.unit
+    def test_omegaconf(self):
+        """Ensure omegaconf raises an error when an unexcepted argument is passed"""
+        with pytest.raises(OmegaConfBaseException):
+            exp_manager(pl.Trainer(), {"unused": 1})
+
+    @pytest.mark.unit
+    def test_trainer_loggers(self, tmp_path):
+        """ Test that a trainer with logger errors out with a number of arguments. Test that it works with
+        create_tensorboard_logger set to False
+        """
+        test_trainer = pl.Trainer()  # Should create logger and modelcheckpoint
+
+        with pytest.raises(LoggerMisconfigurationError):  # Fails because exp_manager defaults to trainer
+            exp_manager(test_trainer, {"exp_dir": str(tmp_path)})
+        with pytest.raises(LoggerMisconfigurationError):  # Fails because exp_manager defaults to trainer
+            exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path)})
+        with pytest.raises(LoggerMisconfigurationError):  # Fails because exp_manager defaults to trainer
+            exp_manager(test_trainer, {"resume_if_exists": True})
+
+        # Check that exp_manager uses trainer.logger, it's exp_dir, name, and version
+        log_dir = exp_manager(test_trainer, {"create_tensorboard_logger": False, "create_checkpoint_callback": False})
+        assert log_dir.resolve() == Path("./lightning_logs/version_0").resolve()
+        assert Path("./lightning_logs").exists()
+        assert Path("./lightning_logs/version_0").exists()
+
+        # Check that a trainer without a logger gets a logger attached to it
+        test_trainer = pl.Trainer(logger=False)
+        log_dir = exp_manager(
+            test_trainer,
+            {"create_tensorboard_logger": True, "create_checkpoint_callback": False, "exp_dir": str(tmp_path)},
+        )
+        assert isinstance(test_trainer.logger, pl.loggers.TensorBoardLogger)
+
+        test_trainer = pl.Trainer(logger=False)
+        # Check that a create_wandb_logger=True errors out unless wandb_logger_kwargs is passed.
+        with pytest.raises(ValueError):
+            log_dir = exp_manager(
+                test_trainer,
+                {
+                    "create_tensorboard_logger": False,
+                    "create_checkpoint_callback": False,
+                    "exp_dir": str(tmp_path),
+                    "create_wandb_logger": True,
+                },
+            )
+        # Check that a WandbLogger is attached to logger if create_wandb_logger=True and wandb_logger_kwargs has name
+        # and project
+        log_dir = exp_manager(
+            test_trainer,
+            {
+                "create_tensorboard_logger": False,
+                "create_checkpoint_callback": False,
+                "exp_dir": str(tmp_path),
+                "create_wandb_logger": True,
+                "wandb_logger_kwargs": {"name": "", "project": ""},
+            },
+        )
+        assert isinstance(test_trainer.logger, pl.loggers.WandbLogger)
+
+    @pytest.mark.unit
+    def test_checkpoint_configurations(self):
+        """ Test that trainer creating modelcheckpoint and asking exp_manager to do it too results in errors, but
+        is error free if only one is asked to do so.
+        """
+        disable_tb_logger = {"create_tensorboard_logger": False}
+        test_trainer = pl.Trainer()  # Should create logger and modelcheckpoint
+        with pytest.raises(CheckpointMisconfigurationError):  # Fails because both try to create modelcheckpoint
+            exp_manager(test_trainer, disable_tb_logger)
+
+        # Should succeed without error
+        exp_manager(test_trainer, {"create_checkpoint_callback": False, "create_tensorboard_logger": False})
+
+        test_trainer_2 = pl.Trainer(checkpoint_callback=False)
+        exp_manager(test_trainer_2, disable_tb_logger)  # Should succeed without error
+
+    @pytest.mark.unit
+    def test_default_log_dir(self):
+        """Check the default of ./nemo_experiments/default/datetime works as intended"""
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
+
+        log_dir = exp_manager(test_trainer, {"create_tensorboard_logger": False, "create_checkpoint_callback": False})
+        assert (log_dir / "..").resolve() == Path("./nemo_experiments/default/").resolve()
+        assert Path("./nemo_experiments").exists()
+        assert Path("./nemo_experiments/default/").exists()
+        sub_dirs = [x for x in Path("./nemo_experiments/default/").iterdir() if x.is_dir()]
+        assert len(sub_dirs) == 1
+        assert re.match(r"[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}-[0-9]{2}-[0-9]{2}", sub_dirs[0].name)
+
+    @pytest.mark.unit
+    def test_log_dir_overrides(self, monkeypatch, tmp_path):
+        """Check a variety of trainer options with exp_manager"""
+        # Checks that explicit_log_dir ignores exp_dir, name, and version
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
+        log_dir = exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path / "test_log_dir_overrides")})
+        assert log_dir.resolve() == (tmp_path / "test_log_dir_overrides").resolve()
+        assert Path(tmp_path).exists()
+        assert Path(tmp_path / "test_log_dir_overrides").exists()
+
+        # Checks that exp_manager uses exp_dir, default name, and explicit version
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
+        log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "version": 957})
+        assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "957").resolve()
+        assert Path(tmp_path).exists()
+        assert Path(tmp_path / "test_no_name" / "default" / "957").exists()
+
+        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
+        # Checks that use_datetime_version False toggle works
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
+        log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "use_datetime_version": False})
+        assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "version_0").resolve()
+        assert Path(tmp_path).exists()
+        assert Path(tmp_path / "test_no_name" / "default" / "version_0").exists()
+
+        monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION)
+        # Checks that use_datetime_version False toggle works and version increments
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
+        log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "use_datetime_version": False})
+        assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "version_1").resolve()
+        assert Path(tmp_path).exists()
+        assert Path(tmp_path / "test_no_name" / "default" / "version_1").exists()
+
+    @pytest.mark.unit
+    def test_resume(self, tmp_path):
+        """ Tests the resume capabilities of exp_manager"""
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
+
+        # Error because explicit_log_dir does not exist
+        with pytest.raises(NotFoundError):
+            exp_manager(
+                test_trainer,
+                {
+                    "exp_dir": str(tmp_path / "test_resume"),
+                    "resume_if_exists": True,
+                    "explicit_log_dir": "Does_not_exist",
+                },
+            )
+
+        # Error because checkpoints folder does not exist
+        with pytest.raises(NotFoundError):
+            exp_manager(test_trainer, {"resume_if_exists": True, "exp_dir": str(tmp_path / "test_resume")})
+
+        # No error because we tell exp_manager to ignore notfounderror
+        exp_manager(
+            test_trainer,
+            {
+                "resume_if_exists": True,
+                "exp_dir": str(tmp_path / "test_resume_2"),
+                "resume_ignore_no_checkpoint": True,
+            },
+        )
+
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
+        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints").mkdir(parents=True)
+        # Error because checkpoints do not exist in folder
+        with pytest.raises(NotFoundError):
+            exp_manager(
+                test_trainer,
+                {
+                    "resume_if_exists": True,
+                    "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"),
+                },
+            )
+
+        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--end.ckpt").touch()
+        # Error because *end.ckpt is in folder indicating that training has already finished
+        with pytest.raises(ValueError):
+            exp_manager(
+                test_trainer,
+                {
+                    "resume_if_exists": True,
+                    "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"),
+                },
+            )
+
+        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--end.ckpt").unlink()
+        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt").touch()
+        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel2--last.ckpt").touch()
+        # Error because multiple *last.ckpt is in folder. If more than one, don't know which to restore
+        with pytest.raises(ValueError):
+            exp_manager(
+                test_trainer,
+                {
+                    "resume_if_exists": True,
+                    "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"),
+                },
+            )
+
+        # Finally succeed
+        Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel2--last.ckpt").unlink()
+        log_dir = exp_manager(
+            test_trainer,
+            {"resume_if_exists": True, "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0")},
+        )
+        checkpoint = Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt")
+        assert Path(test_trainer.resume_from_checkpoint).resolve() == checkpoint.resolve()
+
+        # Succeed again and make sure that run_0 exists and previous log files were moved
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False)
+        exp_manager(test_trainer, {"resume_if_exists": True, "explicit_log_dir": str(log_dir)})
+        checkpoint = Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt")
+        assert Path(test_trainer.resume_from_checkpoint).resolve() == checkpoint.resolve()
+        prev_run_dir = Path(tmp_path / "test_resume" / "default" / "version_0" / "run_0")
+        assert prev_run_dir.exists()
+        prev_log = Path(tmp_path / "test_resume" / "default" / "version_0" / "run_0" / "lightning_logs.txt")
+        assert prev_log.exists()
+
+    @pytest.mark.unit
+    def test_nemo_checkpoint_save_best_model_1(self, tmp_path):
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False, max_epochs=4)
+        log_dir = exp_manager(
+            test_trainer,
+            {"checkpoint_callback_params": {"save_best_model": True}, "explicit_log_dir": str(tmp_path / "test")},
+        )
+        model = ExampleModel()
+        test_trainer.fit(model)
+
+        assert Path(str(tmp_path / "test" / "checkpoints" / "default.nemo")).exists()
+
+        model = ExampleModel.restore_from(str(tmp_path / "test" / "checkpoints" / "default.nemo"))
+        assert float(model(torch.tensor([1.0, 1.0], device=model.device))) == 0.0
+
+    @pytest.mark.unit
+    def test_nemo_checkpoint_save_best_model_2(self, tmp_path):
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False, max_epochs=4)
+        log_dir = exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path / "test")},)
+        model = ExampleModel()
+        test_trainer.fit(model)
+
+        assert Path(str(tmp_path / "test" / "checkpoints" / "default.nemo")).exists()
+
+        model = ExampleModel.restore_from(str(tmp_path / "test" / "checkpoints" / "default.nemo"))
+        assert math.fabs(float(model(torch.tensor([1.0, 1.0], device=model.device))) - 0.03) < 1e-5
+
+    @pytest.mark.unit
+    def test_nemo_checkpoint_always_save_nemo(self, tmp_path):
+        test_trainer = pl.Trainer(checkpoint_callback=False, logger=False, max_epochs=4)
+        log_dir = exp_manager(
+            test_trainer,
+            {
+                "checkpoint_callback_params": {"save_best_model": True, "always_save_nemo": True},
+                "explicit_log_dir": str(tmp_path / "test"),
+            },
+        )
+        model = ExampleModel()
+        test_trainer.fit(model)
+
+        assert Path(str(tmp_path / "test" / "checkpoints" / "default.nemo")).exists()
+
+        model = ExampleModel.restore_from(str(tmp_path / "test" / "checkpoints" / "default.nemo"))
+        assert float(model(torch.tensor([1.0, 1.0], device=model.device))) == 0.0

tts/NeMo/tests/core/test_fileio.py

@@ -0,0 +1,299 @@
+# 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
+import tempfile
+
+import numpy as np
+import pytest
+import torch
+from omegaconf import DictConfig, OmegaConf
+
+from nemo.collections.asr.models import EncDecCTCModel
+
+try:
+    from eff.cookbooks import NeMoCookbook
+
+    _EFF_PRESENT_ = True
+except ImportError:
+    _EFF_PRESENT_ = False
+
+# A decorator marking the EFF requirement.
+requires_eff = pytest.mark.skipif(not _EFF_PRESENT_, reason="Export File Format library required to run test")
+
+
+@pytest.fixture()
+def asr_model():
+    preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+    encoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+        'params': {
+            'feat_in': 64,
+            'activation': 'relu',
+            'conv_mask': True,
+            'jasper': [
+                {
+                    'filters': 1024,
+                    'repeat': 1,
+                    'kernel': [1],
+                    'stride': [1],
+                    'dilation': [1],
+                    'dropout': 0.0,
+                    'residual': False,
+                    'separable': True,
+                    'se': True,
+                    'se_context_size': -1,
+                }
+            ],
+        },
+    }
+
+    decoder = {
+        'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
+        'params': {
+            'feat_in': 1024,
+            'num_classes': 28,
+            'vocabulary': [
+                ' ',
+                'a',
+                'b',
+                'c',
+                'd',
+                'e',
+                'f',
+                'g',
+                'h',
+                'i',
+                'j',
+                'k',
+                'l',
+                'm',
+                'n',
+                'o',
+                'p',
+                'q',
+                'r',
+                's',
+                't',
+                'u',
+                'v',
+                'w',
+                'x',
+                'y',
+                'z',
+                "'",
+            ],
+        },
+    }
+    modelConfig = DictConfig(
+        {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+    )
+
+    model_instance = EncDecCTCModel(cfg=modelConfig)
+    return model_instance
+
+
+class TestFileIO:
+    @pytest.mark.unit
+    def test_to_from_config_file(self, asr_model):
+        """" Test makes sure that the second instance created with the same configuration (BUT NOT checkpoint)
+        has different weights. """
+
+        with tempfile.NamedTemporaryFile() as fp:
+            yaml_filename = fp.name
+            asr_model.to_config_file(path2yaml_file=yaml_filename)
+            next_instance = EncDecCTCModel.from_config_file(path2yaml_file=yaml_filename)
+
+            assert isinstance(next_instance, EncDecCTCModel)
+
+            assert len(next_instance.decoder.vocabulary) == 28
+            assert asr_model.num_weights == next_instance.num_weights
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = next_instance.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert not np.array_equal(w1, w2)
+
+    @pytest.mark.unit
+    def test_save_restore_from_nemo_file(self, asr_model):
+        """" Test makes sure that the second instance created from the same configuration AND checkpoint 
+        has the same weights. """
+
+        with tempfile.NamedTemporaryFile() as fp:
+            filename = fp.name
+
+            # Save model (with random artifact).
+            with tempfile.NamedTemporaryFile() as artifact:
+                asr_model.register_artifact(config_path="abc", src=artifact.name)
+                asr_model.save_to(save_path=filename)
+
+            # Restore the model.
+            asr_model2 = EncDecCTCModel.restore_from(restore_path=filename)
+
+            assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary)
+            assert asr_model.num_weights == asr_model2.num_weights
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert np.array_equal(w1, w2)
+
+    @requires_eff
+    @pytest.mark.unit
+    def test_eff_save_restore_from_nemo_file_encrypted(self, asr_model):
+        """" Test makes sure that after encrypted save-restore the model has the same weights. """
+
+        with tempfile.NamedTemporaryFile() as fp:
+            filename = fp.name
+
+            # Set key - use checkpoint encryption.
+            NeMoCookbook.set_encryption_key("test_key")
+
+            # Save model (with random artifact).
+            with tempfile.NamedTemporaryFile() as artifact:
+                asr_model.register_artifact(config_path="abc", src=artifact.name)
+                asr_model.save_to(save_path=filename)
+
+            # Try to restore the encrypted archive (weights) without the encryption key.
+            NeMoCookbook.set_encryption_key(None)
+            with pytest.raises(PermissionError):
+                # Restore the model.
+                asr_model2 = EncDecCTCModel.restore_from(restore_path=filename)
+
+            # Restore the model.
+            NeMoCookbook.set_encryption_key("test_key")
+            asr_model3 = EncDecCTCModel.restore_from(restore_path=filename)
+            # Reset encryption so it won't mess up with other save/restore.
+            NeMoCookbook.set_encryption_key(None)
+
+            assert asr_model.num_weights == asr_model3.num_weights
+
+    @pytest.mark.unit
+    def test_save_restore_from_nemo_file_with_override(self, asr_model, tmpdir):
+        """" Test makes sure that the second instance created from the same configuration AND checkpoint
+        has the same weights.
+
+        Args:
+            tmpdir: fixture providing a temporary directory unique to the test invocation.
+        """
+        # Name of the archive in tmp folder.
+        filename = os.path.join(tmpdir, "eff.nemo")
+
+        # Get path where the command is executed - the artifacts will be "retrieved" there.
+        # (original .nemo behavior)
+        cwd = os.getcwd()
+
+        with tempfile.NamedTemporaryFile(mode='a+') as conf_fp:
+
+            # Create a "random artifact".
+            with tempfile.NamedTemporaryFile(mode="w", delete=False) as artifact:
+                artifact.write("magic content 42")
+            # Remember the filename of the artifact.
+            _, artifact_filename = os.path.split(artifact.name)
+            # Add artifact to model.
+            asr_model.register_artifact(config_path="abc", src=artifact.name)
+            # Save model (with "random artifact").
+            asr_model.save_to(save_path=filename)
+
+            # Modify config slightly
+            cfg = asr_model.cfg
+            cfg.encoder.activation = 'swish'
+            yaml_cfg = OmegaConf.to_yaml(cfg)
+            conf_fp.write(yaml_cfg)
+            conf_fp.seek(0)
+
+            # Restore the model.
+            asr_model2 = EncDecCTCModel.restore_from(restore_path=filename, override_config_path=conf_fp.name)
+
+            assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary)
+            assert asr_model.num_weights == asr_model2.num_weights
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert np.array_equal(w1, w2)
+
+            assert asr_model2.cfg.encoder.activation == 'swish'
+
+    @pytest.mark.unit
+    def test_save_model_level_pt_ckpt(self, asr_model):
+        with tempfile.TemporaryDirectory() as ckpt_dir:
+            nemo_file = os.path.join(ckpt_dir, 'asr.nemo')
+            asr_model.save_to(nemo_file)
+
+            # Save model level PT checkpoint
+            asr_model.extract_state_dict_from(nemo_file, ckpt_dir)
+            ckpt_path = os.path.join(ckpt_dir, 'model_weights.ckpt')
+
+            assert os.path.exists(ckpt_path)
+
+            # Restore the model.
+            asr_model2 = EncDecCTCModel.restore_from(restore_path=nemo_file)
+
+            assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary)
+            assert asr_model.num_weights == asr_model2.num_weights
+
+            # Change weights values
+            asr_model2.encoder.encoder[0].mconv[0].conv.weight.data += 1.0
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert not np.array_equal(w1, w2)
+
+            # Restore from checkpoint
+            asr_model2.load_state_dict(torch.load(ckpt_path))
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert np.array_equal(w1, w2)
+
+    @pytest.mark.unit
+    def test_save_module_level_pt_ckpt(self, asr_model):
+        with tempfile.TemporaryDirectory() as ckpt_dir:
+            nemo_file = os.path.join(ckpt_dir, 'asr.nemo')
+            asr_model.save_to(nemo_file)
+
+            # Save model level PT checkpoint
+            asr_model.extract_state_dict_from(nemo_file, ckpt_dir, split_by_module=True)
+            encoder_path = os.path.join(ckpt_dir, 'encoder.ckpt')
+            decoder_path = os.path.join(ckpt_dir, 'decoder.ckpt')
+            preprocessor_path = os.path.join(ckpt_dir, 'preprocessor.ckpt')
+
+            assert os.path.exists(encoder_path)
+            assert os.path.exists(decoder_path)
+            assert os.path.exists(preprocessor_path)
+
+            # Restore the model.
+            asr_model2 = EncDecCTCModel.restore_from(restore_path=nemo_file)
+
+            assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary)
+            assert asr_model.num_weights == asr_model2.num_weights
+
+            # Change weights values
+            asr_model2.encoder.encoder[0].mconv[0].conv.weight.data += 1.0
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert not np.array_equal(w1, w2)
+
+            # Restore from checkpoint
+            asr_model2.encoder.load_state_dict(torch.load(encoder_path))
+
+            w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+            w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy()
+
+            assert np.array_equal(w1, w2)

tts/NeMo/tests/core/test_neural_types.py

@@ -0,0 +1,224 @@
+# 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 pytest
+
+from nemo.core.neural_types import (
+    AcousticEncodedRepresentation,
+    AudioSignal,
+    AxisKind,
+    AxisKindAbstract,
+    AxisType,
+    ChannelType,
+    ElementType,
+    MelSpectrogramType,
+    MFCCSpectrogramType,
+    NeuralType,
+    NeuralTypeComparisonResult,
+    SpectrogramType,
+    VoidType,
+)
+
+
+class TestNeuralTypeSystem:
+    @pytest.mark.unit
+    def test_short_vs_long_version(self):
+        long_version = NeuralType(
+            axes=(AxisType(AxisKind.Batch, None), AxisType(AxisKind.Dimension, None), AxisType(AxisKind.Time, None)),
+            elements_type=AcousticEncodedRepresentation(),
+        )
+        short_version = NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation())
+        assert long_version.compare(short_version) == NeuralTypeComparisonResult.SAME
+        assert short_version.compare(long_version) == NeuralTypeComparisonResult.SAME
+
+    @pytest.mark.unit
+    def test_parameterized_type_audio_sampling_frequency(self):
+        audio16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000))
+        audio8K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(8000))
+        another16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000))
+
+        assert audio8K.compare(audio16K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+        assert audio16K.compare(audio8K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS
+        assert another16K.compare(audio16K) == NeuralTypeComparisonResult.SAME
+        assert audio16K.compare(another16K) == NeuralTypeComparisonResult.SAME
+
+    @pytest.mark.unit
+    def test_transpose_same_1(self):
+        type1 = NeuralType(axes=('B', 'T', 'C'))
+        type2 = NeuralType(axes=('T', 'B', 'C'))
+
+        assert type1.compare(type2) == NeuralTypeComparisonResult.TRANSPOSE_SAME
+        assert type2.compare(type1) == NeuralTypeComparisonResult.TRANSPOSE_SAME
+
+    @pytest.mark.unit
+    def test_transpose_same_2(self):
+        audio16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000))
+        audio16K_t = NeuralType(axes=('T', 'B'), elements_type=AudioSignal(16000))
+
+        assert audio16K.compare(audio16K_t) == NeuralTypeComparisonResult.TRANSPOSE_SAME
+
+    @pytest.mark.unit
+    def test_inheritance_spec_augment_example(self):
+        input = NeuralType(('B', 'D', 'T'), SpectrogramType())
+        out1 = NeuralType(('B', 'D', 'T'), MelSpectrogramType())
+        out2 = NeuralType(('B', 'D', 'T'), MFCCSpectrogramType())
+
+        assert out1.compare(out2) == NeuralTypeComparisonResult.INCOMPATIBLE
+        assert out2.compare(out1) == NeuralTypeComparisonResult.INCOMPATIBLE
+        assert input.compare(out1) == NeuralTypeComparisonResult.GREATER
+        assert input.compare(out2) == NeuralTypeComparisonResult.GREATER
+        assert out1.compare(input) == NeuralTypeComparisonResult.LESS
+        assert out2.compare(input) == NeuralTypeComparisonResult.LESS
+
+    @pytest.mark.unit
+    def test_singletone(self):
+        loss_output1 = NeuralType(axes=None)
+        loss_output2 = NeuralType(axes=None)
+
+        assert loss_output1.compare(loss_output2) == NeuralTypeComparisonResult.SAME
+        assert loss_output2.compare(loss_output1) == NeuralTypeComparisonResult.SAME
+
+    @pytest.mark.unit
+    def test_list_of_lists(self):
+        T1 = NeuralType(
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=True),
+                AxisType(kind=AxisKind.Time, size=None, is_list=True),
+                AxisType(kind=AxisKind.Dimension, size=32, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=128, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=256, is_list=False),
+            ),
+            elements_type=ChannelType(),
+        )
+        T2 = NeuralType(
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=False),
+                AxisType(kind=AxisKind.Time, size=None, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=32, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=128, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=256, is_list=False),
+            ),
+            elements_type=ChannelType(),
+        )
+        # TODO: should this be incompatible instead???
+        assert T1.compare(T2), NeuralTypeComparisonResult.TRANSPOSE_SAME
+
+    @pytest.mark.unit
+    def test_void(self):
+        btc_spctr = NeuralType(('B', 'T', 'C'), SpectrogramType())
+        btc_spct_bad = NeuralType(('B', 'T'), SpectrogramType())
+        btc_void = NeuralType(('B', 'T', 'C'), VoidType())
+
+        assert btc_void.compare(btc_spctr) == NeuralTypeComparisonResult.SAME
+        assert btc_spctr.compare(btc_void) == NeuralTypeComparisonResult.INCOMPATIBLE
+        assert btc_void.compare(btc_spct_bad) == NeuralTypeComparisonResult.INCOMPATIBLE
+
+    @pytest.mark.unit
+    def test_big_void(self):
+        big_void_1 = NeuralType(elements_type=VoidType())
+        big_void_2 = NeuralType()
+
+        btc_spctr = NeuralType(('B', 'T', 'C'), SpectrogramType())
+        btc_spct_bad = NeuralType(('B', 'T'), SpectrogramType())
+        t1 = NeuralType(
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=True),
+                AxisType(kind=AxisKind.Time, size=None, is_list=True),
+                AxisType(kind=AxisKind.Dimension, size=32, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=128, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=256, is_list=False),
+            ),
+            elements_type=ChannelType(),
+        )
+        t2 = NeuralType(
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=False),
+                AxisType(kind=AxisKind.Time, size=None, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=32, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=128, is_list=False),
+                AxisType(kind=AxisKind.Dimension, size=256, is_list=False),
+            ),
+            elements_type=ChannelType(),
+        )
+
+        assert big_void_1.compare(btc_spctr) == NeuralTypeComparisonResult.SAME
+        assert big_void_1.compare(btc_spct_bad) == NeuralTypeComparisonResult.SAME
+        assert big_void_1.compare(t1) == NeuralTypeComparisonResult.SAME
+        assert big_void_1.compare(t2) == NeuralTypeComparisonResult.SAME
+
+        assert big_void_2.compare(btc_spctr) == NeuralTypeComparisonResult.SAME
+        assert big_void_2.compare(btc_spct_bad) == NeuralTypeComparisonResult.SAME
+        assert big_void_2.compare(t1) == NeuralTypeComparisonResult.SAME
+        assert big_void_2.compare(t2) == NeuralTypeComparisonResult.SAME
+
+    @pytest.mark.unit
+    def test_unspecified_dimensions(self):
+        t0 = NeuralType(
+            (AxisType(AxisKind.Batch, 64), AxisType(AxisKind.Time, 10), AxisType(AxisKind.Dimension, 128)),
+            SpectrogramType(),
+        )
+        t1 = NeuralType(('B', 'T', 'C'), SpectrogramType())
+
+        assert t1.compare(t0), NeuralTypeComparisonResult.SAME
+        assert t0.compare(t1), NeuralTypeComparisonResult.DIM_INCOMPATIBLE
+
+    @pytest.mark.unit
+    def test_any_axis(self):
+        t0 = NeuralType(('B', 'Any', 'Any'), VoidType())
+        t1 = NeuralType(('B', 'Any', 'Any'), SpectrogramType())
+        t2 = NeuralType(('B', 'T', 'C'), SpectrogramType())
+
+        assert t0.compare(t1) == NeuralTypeComparisonResult.SAME
+        assert t0.compare(t2) == NeuralTypeComparisonResult.SAME
+        assert t1.compare(t2) == NeuralTypeComparisonResult.SAME
+        assert t2.compare(t1) == NeuralTypeComparisonResult.INCOMPATIBLE
+        assert t1.compare(t0) == NeuralTypeComparisonResult.INCOMPATIBLE
+
+    @pytest.mark.unit
+    def test_struct(self):
+        class BoundingBox(ElementType):
+            def __str__(self):
+                return "bounding box from detection model"
+
+            def fields(self):
+                return ("X", "Y", "W", "H")
+
+        # ALSO ADD new, user-defined, axis kind
+        class AxisKind2(AxisKindAbstract):
+            Image = 0
+
+        T1 = NeuralType(
+            elements_type=BoundingBox(),
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=True),
+                AxisType(kind=AxisKind2.Image, size=None, is_list=True),
+            ),
+        )
+
+        class BadBoundingBox(ElementType):
+            def __str__(self):
+                return "bad bounding box from detection model"
+
+            def fields(self):
+                return ("X", "Y", "H")
+
+        T2 = NeuralType(
+            elements_type=BadBoundingBox(),
+            axes=(
+                AxisType(kind=AxisKind.Batch, size=None, is_list=True),
+                AxisType(kind=AxisKind2.Image, size=None, is_list=True),
+            ),
+        )
+        assert T2.compare(T1) == NeuralTypeComparisonResult.INCOMPATIBLE

tts/NeMo/tests/core/test_optimizers_schedulers.py

@@ -0,0 +1,809 @@
+# 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 math
+import random
+
+import omegaconf
+import pytest
+import pytorch_lightning as pl
+import torch
+import torch.optim
+
+from nemo.core import config, optim
+from nemo.core.optim.lr_scheduler import AVAILABLE_SCHEDULERS
+from nemo.core.optim.optimizers import AVAILABLE_OPTIMIZERS
+from nemo.utils import logging
+
+
+class TempModel(torch.nn.Module):
+    def __init__(self):
+        super(TempModel, self).__init__()
+        self.layer = torch.nn.Linear(5, 1)
+
+    def forward(self, x):
+        x = self.layer(x)
+        return x
+
+
+class OptCounter(torch.optim.SGD):
+    def __init__(self, *args, **kwargs):
+        self.count = 0
+        super().__init__(*args, **kwargs)
+
+    def step(self, closure=None):
+        try:
+            self.count += 1
+        except AttributeError:
+            self.count = 1
+        super().step(closure)
+
+
+class RandomDataset(torch.utils.data.Dataset):
+    def __init__(self, dataset_len):
+        super().__init__()
+        self.__dataset_len = dataset_len
+
+    def __getitem__(self, *args):
+        return torch.randn(2)
+
+    def __len__(self):
+        return self.__dataset_len
+
+
+class ExampleModel(pl.LightningModule):
+    def __init__(self, batch_size, dataset_len, drop_last, max_steps):
+        super().__init__()
+        self.l1 = torch.nn.modules.Linear(in_features=2, out_features=1)
+        self.batch_size = batch_size
+        self.dataset_len = dataset_len
+        self.drop_last = drop_last
+        self.max_steps = max_steps
+
+    def train_dataloader(self):
+        dataset = RandomDataset(self.dataset_len)
+        return torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, drop_last=self.drop_last)
+
+    def training_step(self, batch, batch_idx):
+        output = self.l1(batch)
+        output = torch.nn.functional.l1_loss(output, torch.ones(output.size()).to(output.device))
+        return {"loss": output}
+
+    def configure_optimizers(self):
+        self.my_opt = OptCounter(self.parameters(), lr=0.02)
+        return self.my_opt
+
+
+class Callback(pl.callbacks.Callback):
+    @pl.utilities.distributed.rank_zero_only
+    def on_train_end(self, trainer, module):
+        if trainer.global_step != module.my_opt.count or trainer.global_step != module.max_steps:
+            logging.debug(f"max_epochs: {trainer.max_epochs}")
+            logging.debug(f"accumulate_grad_batches: {trainer.accumulate_grad_batches}")
+            logging.debug(f"limit_train_batches: {trainer.limit_train_batches}")
+            logging.debug(f"num_processes: {trainer.num_processes}")
+            logging.debug(f"batch_size: {module.batch_size}")
+            logging.debug(f"dataset_len: {module.dataset_len}")
+            logging.debug(f"drop_last: {module.drop_last}")
+            logging.debug(f"{len(trainer.train_dataloader)}")
+            logging.debug(f"{trainer.num_training_batches }")
+        assert (
+            trainer.global_step == module.my_opt.count
+        ), f"{trainer.global_step} != {module.my_opt.count} != {module.max_steps}"
+        assert (
+            trainer.global_step == module.max_steps
+        ), f"{trainer.global_step} != {module.my_opt.count} != {module.max_steps}"
+
+
+class TestOptimizersSchedulers:
+    INITIAL_LR = 0.1
+    MIN_LR = 1e-3
+    MAX_STEPS = 10
+
+    @pytest.mark.unit
+    def test_get_optimizer(self):
+        model = TempModel()
+
+        for opt_name in AVAILABLE_OPTIMIZERS.keys():
+            opt_cls = optim.get_optimizer(opt_name)
+            opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+            assert isinstance(opt, AVAILABLE_OPTIMIZERS[opt_name])
+
+    @pytest.mark.unit
+    def test_register_optimizer(self):
+        class TempOpt(torch.optim.SGD):
+            pass
+
+        class TempOptParams(config.optimizers.SGDParams):
+            pass
+
+        optim.register_optimizer('TempOpt', TempOpt, TempOptParams)
+
+        model = TempModel()
+        opt_cls = optim.get_optimizer('TempOpt')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        assert isinstance(opt, TempOpt)
+
+    @pytest.mark.unit
+    def test_optim_config_parse_bypass(self):
+        basic_optim_config = {'weight_decay': 0.001, 'betas': [0.8, 0.5]}
+        parsed_params = optim.parse_optimizer_args('novograd', basic_optim_config)
+        assert parsed_params['weight_decay'] == basic_optim_config['weight_decay']
+        assert parsed_params['betas'][0] == basic_optim_config['betas'][0]
+        assert parsed_params['betas'][1] == basic_optim_config['betas'][1]
+
+        dict_config = omegaconf.OmegaConf.create(basic_optim_config)
+        parsed_params = optim.parse_optimizer_args('novograd', dict_config)
+        assert parsed_params['weight_decay'] == dict_config['weight_decay']
+        assert parsed_params['betas'][0] == dict_config['betas'][0]
+        assert parsed_params['betas'][1] == dict_config['betas'][1]
+
+    @pytest.mark.unit
+    def test_optim_config_parse_arg_by_name(self):
+        basic_optim_config = {'name': 'auto', 'weight_decay': 0.001, 'betas': [0.8, 0.5]}
+        parsed_params = optim.parse_optimizer_args('novograd', basic_optim_config)
+        assert parsed_params['weight_decay'] == basic_optim_config['weight_decay']
+        assert parsed_params['betas'][0] == basic_optim_config['betas'][0]
+        assert parsed_params['betas'][1] == basic_optim_config['betas'][1]
+
+        dict_config = omegaconf.OmegaConf.create(basic_optim_config)
+        parsed_params = optim.parse_optimizer_args('novograd', dict_config)
+        assert parsed_params['weight_decay'] == dict_config['weight_decay']
+        assert parsed_params['betas'][0] == dict_config['betas'][0]
+        assert parsed_params['betas'][1] == dict_config['betas'][1]
+
+        with pytest.raises(omegaconf.errors.ConfigKeyError):
+            optim.parse_optimizer_args('sgd', dict_config)
+
+    @pytest.mark.unit
+    def test_optim_config_parse_arg_by_target(self):
+        basic_optim_config = {
+            '_target_': 'nemo.core.config.NovogradParams',
+            'params': {'weight_decay': 0.001, 'betas': [0.8, 0.5]},
+        }
+        basic_optim_config = omegaconf.OmegaConf.create(basic_optim_config)
+        parsed_params = optim.parse_optimizer_args('novograd', basic_optim_config)
+        assert parsed_params['weight_decay'] == basic_optim_config['params']['weight_decay']
+        assert parsed_params['betas'][0] == basic_optim_config['params']['betas'][0]
+        assert parsed_params['betas'][1] == basic_optim_config['params']['betas'][1]
+
+        dict_config = omegaconf.OmegaConf.create(basic_optim_config)
+        parsed_params = optim.parse_optimizer_args('novograd', dict_config)
+        assert parsed_params['weight_decay'] == dict_config['params']['weight_decay']
+        assert parsed_params['betas'][0] == dict_config['params']['betas'][0]
+        assert parsed_params['betas'][1] == dict_config['params']['betas'][1]
+
+        # Names are ignored when passing class path
+        # This will be captured during optimizer instantiation
+        output_config = optim.parse_optimizer_args('sgd', dict_config)
+        sgd_config = vars(config.SGDParams())
+        novograd_config = vars(config.NovogradParams())
+
+        assert set(output_config.keys()) != set(sgd_config.keys())
+        assert set(output_config.keys()) == set(novograd_config)
+
+    @pytest.mark.unit
+    def test_get_scheduler(self):
+        model = TempModel()
+        optimizer = optim.Novograd(model.parameters(), lr=self.INITIAL_LR)
+
+        for sched_name in AVAILABLE_SCHEDULERS.keys():
+            sched_cls = optim.lr_scheduler.get_scheduler(sched_name)
+
+            try:
+                sched = sched_cls(optimizer)
+                assert isinstance(sched, AVAILABLE_SCHEDULERS[sched_name])
+                continue
+            except Exception:
+                pass
+
+            try:
+                sched = sched_cls(optimizer, max_steps=self.MAX_STEPS)
+                assert isinstance(sched, AVAILABLE_SCHEDULERS[sched_name])
+                continue
+            except Exception:
+                pass
+
+    @pytest.mark.unit
+    def test_register_scheduler(self):
+        class TempSched(optim.lr_scheduler.CosineAnnealing):
+            pass
+
+        class TempSchedParams(config.schedulers.CosineAnnealingParams):
+            pass
+
+        optim.lr_scheduler.register_scheduler('TempSched', TempSched, TempSchedParams)
+
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+        sched_cls = optim.lr_scheduler.get_scheduler('TempSched')
+        sched = sched_cls(opt, max_steps=self.MAX_STEPS)
+
+        assert isinstance(sched, TempSched)
+
+    @pytest.mark.unit
+    def test_sched_config_parse_simple(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        basic_sched_config = {'name': 'CosineAnnealing', 'max_steps': 10}
+        scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, basic_sched_config)
+        assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing)
+
+        dict_config = omegaconf.OmegaConf.create(basic_sched_config)
+        scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, dict_config)
+        assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing)
+
+    @pytest.mark.unit
+    def test_sched_config_parse_from_cls(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        basic_sched_config = {
+            '_target_': 'nemo.core.config.CosineAnnealingParams',
+            'params': {'min_lr': 0.1},
+            'max_steps': self.MAX_STEPS,
+        }
+        scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, basic_sched_config)
+        assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing)
+
+        dict_config = omegaconf.OmegaConf.create(basic_sched_config)
+        scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, dict_config)
+        assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing)
+
+    @pytest.mark.unit
+    def test_WarmupPolicy(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        # No warmup case
+        policy = optim.lr_scheduler.WarmupPolicy(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr == self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            assert policy.get_last_lr()[0] == self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup steps available
+        policy = optim.lr_scheduler.WarmupPolicy(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 4:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] == self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+    @pytest.mark.unit
+    def test_WarmupHoldPolicy(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        # No warmup case
+        policy = optim.lr_scheduler.WarmupHoldPolicy(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr == self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            assert policy.get_last_lr()[0] == self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup steps available
+        policy = optim.lr_scheduler.WarmupHoldPolicy(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 4:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] == self.INITIAL_LR
+
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup + Hold steps available
+        policy = optim.lr_scheduler.WarmupHoldPolicy(
+            opt, warmup_steps=5, hold_steps=3, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR
+        )
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 4:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] == self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+    @pytest.mark.unit
+    def test_WarmupAnnealing(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        # No warmup case
+        policy = optim.lr_scheduler.WarmupAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr == self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup steps available
+        policy = optim.lr_scheduler.WarmupAnnealing(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 5:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] < self.INITIAL_LR
+
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup + Hold steps available
+        policy = optim.lr_scheduler.WarmupHoldPolicy(
+            opt, warmup_steps=5, hold_steps=3, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR
+        )
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 4:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] == self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+    @pytest.mark.unit
+    def test_SquareAnnealing(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        # No warmup case
+        policy = optim.lr_scheduler.SquareAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr == self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup steps available
+        policy = optim.lr_scheduler.SquareAnnealing(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 5:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] < self.INITIAL_LR
+
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+    @pytest.mark.unit
+    def test_SquareRootAnnealing(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        # No warmup case
+        policy = optim.lr_scheduler.SquareRootAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr == self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup steps available
+        policy = optim.lr_scheduler.SquareRootAnnealing(
+            opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR
+        )
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 5:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] < self.INITIAL_LR
+
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+    @pytest.mark.unit
+    def test_CosineAnnealing(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        # No warmup case
+        policy = optim.lr_scheduler.CosineAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr == self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup steps available
+        policy = optim.lr_scheduler.CosineAnnealing(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 5:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] < self.INITIAL_LR
+
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+    @pytest.mark.unit
+    def test_PolynomialDecayAnnealing(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        # No warmup case
+        policy = optim.lr_scheduler.PolynomialDecayAnnealing(
+            opt, power=2, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR
+        )
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr == self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup steps available
+        policy = optim.lr_scheduler.PolynomialDecayAnnealing(
+            opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR
+        )
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 5:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] < self.INITIAL_LR
+
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+    @pytest.mark.unit
+    def test_PolynomialHoldDecayAnnealing(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        # No warmup case
+        policy = optim.lr_scheduler.PolynomialHoldDecayAnnealing(
+            opt, power=2, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR
+        )
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr == self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup steps available
+        policy = optim.lr_scheduler.PolynomialHoldDecayAnnealing(
+            opt, power=2, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR
+        )
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 5:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] < self.INITIAL_LR
+
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup + Hold steps available
+        policy = optim.lr_scheduler.PolynomialHoldDecayAnnealing(
+            opt, warmup_steps=5, hold_steps=3, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR, power=2
+        )
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 4:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            elif i <= 8:
+                assert policy.get_last_lr()[0] == self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] < self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+    @pytest.mark.unit
+    def test_InverseSquareRootAnnealing(self):
+        model = TempModel()
+        opt_cls = optim.get_optimizer('novograd')
+        opt = opt_cls(model.parameters(), lr=self.INITIAL_LR)
+
+        # No warmup case
+        policy = optim.lr_scheduler.InverseSquareRootAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR)
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr == self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+        # Warmup steps available
+        policy = optim.lr_scheduler.InverseSquareRootAnnealing(
+            opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR
+        )
+        initial_lr = policy.get_last_lr()[0]
+
+        assert initial_lr < self.INITIAL_LR
+
+        for i in range(self.MAX_STEPS):
+            if i <= 5:
+                assert policy.get_last_lr()[0] <= self.INITIAL_LR
+            else:
+                assert policy.get_last_lr()[0] < self.INITIAL_LR
+
+            opt.step()
+            policy.step()
+
+        policy.step()
+        final_lr = policy.get_last_lr()[0]
+
+        assert final_lr == self.MIN_LR
+
+    @pytest.mark.unit
+    @pytest.mark.run_only_on('CPU')
+    def test_max_step_computation(self):
+        def train(
+            max_epochs, accumulate_grad_batches, limit_train_batches, num_processes, batch_size, dataset_len, drop_last
+        ):
+            trainer = pl.Trainer(
+                max_epochs=max_epochs,
+                accelerator="ddp_cpu",
+                num_processes=num_processes,
+                accumulate_grad_batches=accumulate_grad_batches,
+                limit_train_batches=limit_train_batches,
+                checkpoint_callback=False,
+                progress_bar_refresh_rate=0,
+                weights_summary=None,
+            )
+            max_steps = optim.lr_scheduler.compute_max_steps(
+                max_epochs,
+                accumulate_grad_batches,
+                limit_train_batches,
+                num_processes,
+                dataset_len,
+                batch_size,
+                drop_last,
+            )
+            model = ExampleModel(batch_size, dataset_len, drop_last, max_steps)
+            trainer.callbacks.append(Callback())
+            trainer.fit(model)
+
+        # This test will break once we and lightning upgrade to pytorch 1.7.0 due to a bug fix in pytorch 1.7.0
+        train(
+            31,
+            accumulate_grad_batches=1,
+            limit_train_batches=1.0,
+            num_processes=9,
+            batch_size=60,
+            dataset_len=1613,
+            drop_last=True,
+        )
+        train(
+            5,
+            accumulate_grad_batches=1,
+            limit_train_batches=0.17382691901706027,
+            num_processes=4,
+            batch_size=97,
+            dataset_len=498,
+            drop_last=False,
+        )
+        train(
+            5,
+            accumulate_grad_batches=8,
+            limit_train_batches=0.1663306588594945,
+            num_processes=4,
+            batch_size=54,
+            dataset_len=629,
+            drop_last=True,
+        )
+        train(
+            5,
+            accumulate_grad_batches=1,
+            limit_train_batches=0.2121376533631948,
+            num_processes=1,
+            batch_size=68,
+            dataset_len=488,
+            drop_last=False,
+        )
+        for _ in range(5):
+            drop_last = bool(random.randint(0, 1))
+            accumulate_grad_batches = random.randint(1, 10)
+
+            limit_train_batches_int = random.randint(1, 10)
+            limit_train_batches_float = random.uniform(0, 1)
+            limit_train_batches = random.choice([limit_train_batches_int, limit_train_batches_float])
+            max_epochs = random.randint(4, 20)
+            num_processes = random.randint(1, 5)
+            dataset_len = random.randint(20, num_processes * 500)
+            batch_size = random.randint(math.ceil(5.0 / num_processes), min(dataset_len // num_processes, 128))
+            train(
+                max_epochs,
+                accumulate_grad_batches,
+                limit_train_batches,
+                num_processes,
+                batch_size,
+                dataset_len,
+                drop_last,
+            )

tts/NeMo/tests/core/test_save_restore.py

@@ -0,0 +1,443 @@
+# 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 filecmp
+import os
+import shutil
+import tempfile
+from typing import Dict, Optional, Set, Union
+
+import pytest
+import torch
+from omegaconf import DictConfig, OmegaConf, open_dict
+
+from nemo.collections.asr.models import EncDecCTCModel, EncDecCTCModelBPE
+from nemo.collections.nlp.models import PunctuationCapitalizationModel, TransformerLMModel
+from nemo.core.classes import ModelPT
+from nemo.utils.app_state import AppState
+
+
+def getattr2(object, attr):
+    if not '.' in attr:
+        return getattr(object, attr)
+    else:
+        arr = attr.split('.')
+        return getattr2(getattr(object, arr[0]), '.'.join(arr[1:]))
+
+
+class MockModel(ModelPT):
+    def __init__(self, cfg, trainer=None):
+        super(MockModel, self).__init__(cfg=cfg, trainer=trainer)
+        self.w = torch.nn.Linear(10, 1)
+        # mock temp file
+        if 'temp_file' in self.cfg and self.cfg.temp_file is not None:
+            self.temp_file = self.register_artifact('temp_file', self.cfg.temp_file)
+            with open(self.temp_file, 'r') as f:
+                self.temp_data = f.readlines()
+        else:
+            self.temp_file = None
+            self.temp_data = None
+
+    def forward(self, x):
+        y = self.w(x)
+        return y, self.cfg.temp_file
+
+    def setup_training_data(self, train_data_config: Union[DictConfig, Dict]):
+        self._train_dl = None
+
+    def setup_validation_data(self, val_data_config: Union[DictConfig, Dict]):
+        self._validation_dl = None
+
+    def setup_test_data(self, test_data_config: Union[DictConfig, Dict]):
+        self._test_dl = None
+
+    def list_available_models(cls):
+        return []
+
+
+def _mock_model_config():
+    conf = {'temp_file': None}
+    conf = OmegaConf.create({'model': conf})
+    OmegaConf.set_struct(conf, True)
+    return conf
+
+
+class TestSaveRestore:
+    def __test_restore_elsewhere(
+        self,
+        model: ModelPT,
+        attr_for_eq_check: Set[str] = None,
+        override_config_path: Optional[Union[str, DictConfig]] = None,
+        map_location: Optional[torch.device] = None,
+        strict: bool = False,
+        return_config: bool = False,
+    ):
+        """Test's logic:
+            1. Save model into temporary folder (save_folder)
+            2. Copy .nemo file from save_folder to restore_folder
+            3. Delete save_folder
+            4. Attempt to restore from .nemo file in restore_folder and compare to original instance
+        """
+        # Create a new temporary directory
+        with tempfile.TemporaryDirectory() as restore_folder:
+            with tempfile.TemporaryDirectory() as save_folder:
+                save_folder_path = save_folder
+                # Where model will be saved
+                model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo")
+                model.save_to(save_path=model_save_path)
+                # Where model will be restored from
+                model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo")
+                shutil.copy(model_save_path, model_restore_path)
+            # at this point save_folder should not exist
+            assert save_folder_path is not None and not os.path.exists(save_folder_path)
+            assert not os.path.exists(model_save_path)
+            assert os.path.exists(model_restore_path)
+            # attempt to restore
+            model_copy = model.__class__.restore_from(
+                restore_path=model_restore_path,
+                map_location=map_location,
+                strict=strict,
+                return_config=return_config,
+                override_config_path=override_config_path,
+            )
+
+            if return_config:
+                return model_copy
+
+            assert model.num_weights == model_copy.num_weights
+            if attr_for_eq_check is not None and len(attr_for_eq_check) > 0:
+                for attr in attr_for_eq_check:
+                    assert getattr2(model, attr) == getattr2(model_copy, attr)
+
+            return model_copy
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_EncDecCTCModel(self):
+        # TODO: Switch to using named configs because here we don't really care about weights
+        qn = EncDecCTCModel.from_pretrained(model_name="QuartzNet15x5Base-En")
+        self.__test_restore_elsewhere(model=qn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"]))
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_EncDecCTCModelBPE(self):
+        # TODO: Switch to using named configs because here we don't really care about weights
+        cn = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_citrinet_256")
+        self.__test_restore_elsewhere(model=cn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"]))
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_EncDecCTCModelBPE_v2(self):
+        # TODO: Switch to using named configs because here we don't really care about weights
+        cn = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_conformer_ctc_small")
+        self.__test_restore_elsewhere(model=cn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"]))
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_PunctuationCapitalization(self):
+        # TODO: Switch to using named configs because here we don't really care about weights
+        pn = PunctuationCapitalizationModel.from_pretrained(model_name='punctuation_en_distilbert')
+        self.__test_restore_elsewhere(
+            model=pn, attr_for_eq_check=set(["punct_classifier.log_softmax", "punct_classifier.log_softmax"])
+        )
+
+    @pytest.mark.unit
+    def test_mock_save_to_restore_from(self):
+        with tempfile.NamedTemporaryFile('w') as empty_file:
+            # Write some data
+            empty_file.writelines(["*****\n"])
+            empty_file.flush()
+
+            # Update config
+            cfg = _mock_model_config()
+            cfg.model.temp_file = empty_file.name
+
+            # Create model
+            model = MockModel(cfg=cfg.model, trainer=None)
+            model = model.to('cpu')
+
+            assert model.temp_file == empty_file.name
+
+            # Save test
+            model_copy = self.__test_restore_elsewhere(model, map_location='cpu')
+
+        # Restore test
+        diff = model.w.weight - model_copy.w.weight
+        # because of caching - cache gets prepended
+        assert os.path.basename(model_copy.temp_file).endswith(os.path.basename(model.temp_file))
+        assert diff.mean() <= 1e-9
+        # assert os.path.basename(model.temp_file) == model_copy.temp_file
+        assert model_copy.temp_data == ["*****\n"]
+
+    @pytest.mark.unit
+    def test_mock_restore_from_config_only(self):
+        with tempfile.NamedTemporaryFile('w') as empty_file:
+            # Write some data
+            empty_file.writelines(["*****\n"])
+            empty_file.flush()
+
+            # Update config
+            cfg = _mock_model_config()
+            cfg.model.temp_file = os.path.abspath(empty_file.name)
+
+            # Inject arbitrary config arguments (after creating model)
+            with open_dict(cfg.model):
+                cfg.model.xyz = "abc"
+
+            # Create model
+            model = MockModel(cfg=cfg.model, trainer=None)
+            model = model.to('cpu')
+
+            assert model.temp_file == empty_file.name
+            model_copy = self.__test_restore_elsewhere(model, map_location='cpu', return_config=False)
+            # because of caching - cache gets prepended
+            assert os.path.basename(model_copy.temp_file).endswith(os.path.basename(model.temp_file))
+            # assert filecmp.cmp(model.temp_file, model_copy._cfg.temp_file)
+            assert model.cfg.xyz == model_copy.cfg.xyz
+
+    @pytest.mark.unit
+    def test_mock_restore_from_config_override_with_OmegaConf(self):
+        with tempfile.NamedTemporaryFile('w') as empty_file:
+            # Write some data
+            empty_file.writelines(["*****\n"])
+            empty_file.flush()
+
+            # Update config
+            cfg = _mock_model_config()
+            cfg.model.temp_file = empty_file.name
+
+            # Create model
+            model = MockModel(cfg=cfg.model, trainer=None)
+            model = model.to('cpu')
+
+            assert model.temp_file == empty_file.name
+
+            # Inject arbitrary config arguments (after creating model)
+            with open_dict(cfg.model):
+                cfg.model.xyz = "abc"
+
+            # Save test (with overriden config as OmegaConf object)
+            model_copy = self.__test_restore_elsewhere(model, map_location='cpu', override_config_path=cfg)
+
+        # Restore test
+        diff = model.w.weight - model_copy.w.weight
+        assert diff.mean() <= 1e-9
+        assert model_copy.temp_data == ["*****\n"]
+
+        # Test that new config has arbitrary content
+        assert model_copy.cfg.xyz == "abc"
+
+    @pytest.mark.unit
+    def test_mock_restore_from_config_override_with_yaml(self):
+        with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as config_file:
+            # Write some data
+            empty_file.writelines(["*****\n"])
+            empty_file.flush()
+
+            # Update config
+            cfg = _mock_model_config()
+            cfg.model.temp_file = empty_file.name
+
+            # Create model
+            model = MockModel(cfg=cfg.model, trainer=None)
+            model = model.to('cpu')
+
+            assert model.temp_file == empty_file.name
+
+            # Inject arbitrary config arguments (after creating model)
+            with open_dict(cfg.model):
+                cfg.model.xyz = "abc"
+
+            # Write new config into file
+            OmegaConf.save(cfg, config_file)
+
+            # Save test (with overriden config as OmegaConf object)
+            model_copy = self.__test_restore_elsewhere(
+                model, map_location='cpu', override_config_path=config_file.name
+            )
+
+            # Restore test
+            diff = model.w.weight - model_copy.w.weight
+            assert diff.mean() <= 1e-9
+            assert filecmp.cmp(model.temp_file, model_copy.temp_file)
+            assert model_copy.temp_data == ["*****\n"]
+
+            # Test that new config has arbitrary content
+            assert model_copy.cfg.xyz == "abc"
+
+    @pytest.mark.unit
+    def test_mock_save_to_restore_from_with_target_class(self):
+        with tempfile.NamedTemporaryFile('w') as empty_file:
+            # Write some data
+            empty_file.writelines(["*****\n"])
+            empty_file.flush()
+
+            # Update config
+            cfg = _mock_model_config()
+            cfg.model.temp_file = empty_file.name
+
+            # Create model
+            model = MockModel(cfg=cfg.model, trainer=None)
+            model = model.to('cpu')  # type: MockModel
+
+            assert model.temp_file == empty_file.name
+
+            # Save file using MockModel
+            with tempfile.TemporaryDirectory() as save_folder:
+                save_path = os.path.join(save_folder, "temp.nemo")
+                model.save_to(save_path)
+
+                # Restore test (using ModelPT as restorer)
+                # This forces the target class = MockModel to be used as resolver
+                model_copy = ModelPT.restore_from(save_path, map_location='cpu')
+            # because of caching - cache gets prepended
+            assert os.path.basename(model_copy.temp_file).endswith(os.path.basename(model.temp_file))
+            # assert filecmp.cmp(model.temp_file, model_copy.temp_file)
+        # Restore test
+        diff = model.w.weight - model_copy.w.weight
+        assert diff.mean() <= 1e-9
+        assert isinstance(model_copy, MockModel)
+        assert model_copy.temp_data == ["*****\n"]
+
+    @pytest.mark.unit
+    def test_mock_save_to_restore_from_multiple_models(self):
+        with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as empty_file2:
+            # Write some data
+            empty_file.writelines(["*****\n"])
+            empty_file.flush()
+            empty_file2.writelines(["+++++\n"])
+            empty_file2.flush()
+
+            # Update config + create ,pde;s
+            cfg = _mock_model_config()
+            cfg.model.temp_file = empty_file.name
+            cfg2 = _mock_model_config()
+            cfg2.model.temp_file = empty_file2.name
+
+            # Create models
+            model = MockModel(cfg=cfg.model, trainer=None)
+            model = model.to('cpu')
+            model2 = MockModel(cfg=cfg2.model, trainer=None)
+            model2 = model2.to('cpu')
+
+            assert model.temp_file == empty_file.name
+            assert model2.temp_file == empty_file2.name
+
+            # Save test
+            model_copy = self.__test_restore_elsewhere(model, map_location='cpu')
+            model2_copy = self.__test_restore_elsewhere(model2, map_location='cpu')
+
+        # Restore test
+        assert model_copy.temp_data == ["*****\n"]
+        assert model2_copy.temp_data == ["+++++\n"]
+
+    @pytest.mark.unit
+    def test_mock_save_to_restore_from_multiple_models_inverted_order(self):
+        with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as empty_file2:
+            # Write some data
+            empty_file.writelines(["*****\n"])
+            empty_file.flush()
+            empty_file2.writelines(["+++++\n"])
+            empty_file2.flush()
+
+            # Update config + create ,pde;s
+            cfg = _mock_model_config()
+            cfg.model.temp_file = empty_file.name
+            cfg2 = _mock_model_config()
+            cfg2.model.temp_file = empty_file2.name
+
+            # Create models
+            model = MockModel(cfg=cfg.model, trainer=None)
+            model = model.to('cpu')
+            model2 = MockModel(cfg=cfg2.model, trainer=None)
+            model2 = model2.to('cpu')
+
+            assert model.temp_file == empty_file.name
+            assert model2.temp_file == empty_file2.name
+
+            # Save test (inverted order)
+            model2_copy = self.__test_restore_elsewhere(model2, map_location='cpu')
+            model_copy = self.__test_restore_elsewhere(model, map_location='cpu')
+
+        # Restore test
+        assert model_copy.temp_data == ["*****\n"]
+        assert model2_copy.temp_data == ["+++++\n"]
+
+    @pytest.mark.unit
+    def test_mock_save_to_restore_chained(self):
+        with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as empty_file2:
+            # Write some data
+            empty_file.writelines(["*****\n"])
+            empty_file.flush()
+
+            # Update config + create ,pde;s
+            cfg = _mock_model_config()
+            cfg.model.temp_file = empty_file.name
+
+            # Create models
+            model = MockModel(cfg=cfg.model, trainer=None)
+            model = model.to('cpu')
+
+            assert model.temp_file == empty_file.name
+
+            def save_copy(model, save_folder, restore_folder):
+                # Where model will be saved
+                model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo")
+                model.save_to(save_path=model_save_path)
+                # Where model will be restored from
+                model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo")
+                shutil.copy(model_save_path, model_restore_path)
+                return model_restore_path
+
+            # Save test
+            with tempfile.TemporaryDirectory() as level4:
+                with tempfile.TemporaryDirectory() as level3:
+                    with tempfile.TemporaryDirectory() as level2:
+                        with tempfile.TemporaryDirectory() as level1:
+                            path = save_copy(model, level1, level2)
+                        model_copy2 = model.__class__.restore_from(path)
+                        path = save_copy(model_copy2, level2, level3)
+                    model_copy3 = model.__class__.restore_from(path)
+                    path = save_copy(model_copy3, level3, level4)
+                model_copy = model.__class__.restore_from(path)
+
+        # Restore test
+        assert model_copy.temp_data == ["*****\n"]
+
+        # AppState test
+        appstate = AppState()
+        metadata = appstate.get_model_metadata_from_guid(model_copy.model_guid)
+        assert metadata.guid != model.model_guid
+        assert metadata.restoration_path == path
+
+    @pytest.mark.unit
+    def test_mock_save_to_multiple_times(self):
+        with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.TemporaryDirectory() as tmpdir:
+            # Write some data
+            empty_file.writelines(["*****\n"])
+            empty_file.flush()
+
+            # Update config
+            cfg = _mock_model_config()
+            cfg.model.temp_file = empty_file.name
+
+            # Create model
+            model = MockModel(cfg=cfg.model, trainer=None)  # type: MockModel
+            model = model.to('cpu')
+
+            assert model.temp_file == empty_file.name
+
+            # Save test
+            model.save_to(os.path.join(tmpdir, 'save_0.nemo'))
+            model.save_to(os.path.join(tmpdir, 'save_1.nemo'))
+            model.save_to(os.path.join(tmpdir, 'save_2.nemo'))

tts/NeMo/tests/core/test_serialization.py

@@ -0,0 +1,165 @@
+# 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 pytest
+from omegaconf import DictConfig
+
+from nemo.collections.asr.models import EncDecCTCModel
+from nemo.collections.asr.modules import SpectrogramAugmentation
+from nemo.core.classes.common import Serialization
+
+
+def get_class_path(cls):
+    return f"{cls.__module__}.{cls.__name__}"
+
+
+class MockSerializationImpl(Serialization):
+    def __init__(self, cfg: DictConfig):
+        self.cfg = cfg
+        self.value = self.__class__.__name__
+
+
+class MockSerializationImplV2(MockSerializationImpl):
+    pass
+
+
+class TestSerialization:
+    @pytest.mark.unit
+    def test_from_config_dict_with_cls(self):
+        """Here we test that instantiation works for configs with cls class path in them.
+        Note that just Serialization.from_config_dict can be used to create an object"""
+        config = DictConfig(
+            {
+                'cls': 'nemo.collections.asr.modules.SpectrogramAugmentation',
+                'params': {'rect_freq': 50, 'rect_masks': 5, 'rect_time': 120,},
+            }
+        )
+        obj = Serialization.from_config_dict(config=config)
+        assert isinstance(obj, SpectrogramAugmentation)
+
+    @pytest.mark.unit
+    def test_from_config_dict_without_cls(self):
+        """Here we test that instantiation works for configs without cls class path in them.
+        IMPORTANT: in this case, correct class type should call from_config_dict. This should work for Models."""
+        preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})}
+        encoder = {
+            'cls': 'nemo.collections.asr.modules.ConvASREncoder',
+            'params': {
+                'feat_in': 64,
+                'activation': 'relu',
+                'conv_mask': True,
+                'jasper': [
+                    {
+                        'filters': 1024,
+                        'repeat': 1,
+                        'kernel': [1],
+                        'stride': [1],
+                        'dilation': [1],
+                        'dropout': 0.0,
+                        'residual': False,
+                        'separable': True,
+                        'se': True,
+                        'se_context_size': -1,
+                    }
+                ],
+            },
+        }
+
+        decoder = {
+            'cls': 'nemo.collections.asr.modules.ConvASRDecoder',
+            'params': {
+                'feat_in': 1024,
+                'num_classes': 28,
+                'vocabulary': [
+                    ' ',
+                    'a',
+                    'b',
+                    'c',
+                    'd',
+                    'e',
+                    'f',
+                    'g',
+                    'h',
+                    'i',
+                    'j',
+                    'k',
+                    'l',
+                    'm',
+                    'n',
+                    'o',
+                    'p',
+                    'q',
+                    'r',
+                    's',
+                    't',
+                    'u',
+                    'v',
+                    'w',
+                    'x',
+                    'y',
+                    'z',
+                    "'",
+                ],
+            },
+        }
+        modelConfig = DictConfig(
+            {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)}
+        )
+        obj = EncDecCTCModel.from_config_dict(config=modelConfig)
+        assert isinstance(obj, EncDecCTCModel)
+
+    @pytest.mark.unit
+    def test_config_updated(self):
+        config = DictConfig(
+            {
+                'cls': 'nemo.collections.asr.modules.SpectrogramAugmentation',
+                'params': {'rect_freq': 50, 'rect_masks': 5, 'rect_time': 120,},
+            }
+        )
+        obj = Serialization.from_config_dict(config=config)
+        new_config = obj.to_config_dict()
+        assert config != new_config
+        assert 'params' not in new_config
+        assert 'cls' not in new_config
+        assert '_target_' in new_config
+
+    @pytest.mark.unit
+    def test_base_class_instantiation(self):
+        # Target class is V2 impl, calling class is Serialization (base class)
+        config = DictConfig({'target': get_class_path(MockSerializationImplV2)})
+        obj = Serialization.from_config_dict(config=config)
+        new_config = obj.to_config_dict()
+        assert config == new_config
+        assert isinstance(obj, MockSerializationImplV2)
+        assert obj.value == "MockSerializationImplV2"
+
+    @pytest.mark.unit
+    def test_self_class_instantiation(self):
+        # Target class is V1 impl, calling class is V1 (same class)
+        config = DictConfig({'target': get_class_path(MockSerializationImpl)})
+        obj = MockSerializationImpl.from_config_dict(config=config)  # Serialization is base class
+        new_config = obj.to_config_dict()
+        assert config == new_config
+        assert isinstance(obj, MockSerializationImpl)
+        assert obj.value == "MockSerializationImpl"
+
+    @pytest.mark.unit
+    def test_sub_class_instantiation(self):
+        # Target class is V1 impl, calling class is V2 (sub class)
+        config = DictConfig({'target': get_class_path(MockSerializationImpl)})
+        obj = MockSerializationImplV2.from_config_dict(config=config)  # Serialization is base class
+        new_config = obj.to_config_dict()
+        assert config == new_config
+        assert isinstance(obj, MockSerializationImplV2)
+        assert obj.value == "MockSerializationImplV2"