Introduce code from Rust monorepo. This is a mess

tts/nemo/collections/cv/__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.cv import datasets, losses, models, modules
+from nemo.package_info import __version__
+
+# Set collection version equal to NeMo version.
+__version = __version__
+
+# Authorship.
+__author__ = "NVIDIA Corporation"
+
+# Set collection name.
+__description__ = "Computer Vision collection"

tts/nemo/collections/cv/datasets/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.cv.datasets.mnist_dataset import MNISTDataset, MNISTDatasetConfig

tts/nemo/collections/cv/datasets/mnist_dataset.py

@@ -0,0 +1,134 @@
+# -*- coding: utf-8 -*-
+# 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 os.path import expanduser
+from typing import Optional
+
+from torchvision.datasets import MNIST
+from torchvision.transforms import Compose, Resize, ToTensor
+
+from nemo.core.classes import Dataset
+from nemo.core.config import Config
+from nemo.core.neural_types.axes import AxisKind, AxisType
+from nemo.core.neural_types.elements import ClassificationTarget, Index, NormalizedImageValue, StringLabel
+from nemo.core.neural_types.neural_type import NeuralType
+from nemo.utils.decorators import add_port_docs, experimental
+
+
+@dataclass
+class MNISTDatasetConfig(Config):
+    """
+    Structured config for MNISTDataset class.
+
+    Args:
+        height: image height (DEFAULT: 28)
+        width: image width (DEFAULT: 28)
+        data_folder: path to the folder with data, can be relative to user (DEFAULT: "~/data/mnist")
+        train: use train or test splits (DEFAULT: True)
+        name: Name of the module (DEFAULT: None)
+    """
+
+    height: int = 28
+    width: int = 28
+    data_folder: str = "~/data/mnist"
+    train: bool = True
+    download: bool = True
+
+
+class MNISTDataset(Dataset):
+    """
+    A "thin wrapper" around the torchvision's MNIST dataset.
+    """
+
+    def __init__(self, cfg: MNISTDatasetConfig = MNISTDatasetConfig()):
+        """
+        Initializes the MNIST dataset.
+
+        Args:
+            cfg: Configuration object of type MNISTDatasetConfig.
+
+        """
+        # Call the base class constructor of Dataset.
+        Dataset.__init__(self)  # , name=name)
+
+        # Store height and width.
+        self._height = cfg.height
+        self._width = cfg.width
+
+        # Create transformations: up-scale and transform to tensors.
+        mnist_transforms = Compose([Resize((self._height, self._width)), ToTensor()])
+
+        # Get absolute path.
+        abs_data_folder = expanduser(cfg.data_folder)
+
+        # Create the MNIST dataset object.
+        self._dataset = MNIST(root=abs_data_folder, train=cfg.train, download=cfg.download, transform=mnist_transforms)
+
+        # Class names.
+        labels = 'Zero One Two Three Four Five Six Seven Eight Nine'.split(' ')
+        word_to_ix = {labels[i]: i for i in range(10)}
+
+        # Reverse mapping.
+        self._ix_to_word = {value: key for (key, value) in word_to_ix.items()}
+
+    @property
+    def output_types(self):
+        """
+        Creates definitions of output ports.
+        """
+        return {
+            "indices": NeuralType(tuple('B'), elements_type=Index()),
+            "images": NeuralType(
+                axes=(
+                    AxisType(kind=AxisKind.Batch),
+                    AxisType(kind=AxisKind.Channel, size=1),
+                    AxisType(kind=AxisKind.Height, size=self._height),
+                    AxisType(kind=AxisKind.Width, size=self._width),
+                ),
+                elements_type=NormalizedImageValue(),  # float, <0-1>
+            ),
+            "targets": NeuralType(tuple('B'), elements_type=ClassificationTarget()),  # Target are ints!
+            "labels": NeuralType(tuple('B'), elements_type=StringLabel()),  # Labels is string!
+        }
+
+    def __len__(self):
+        """
+        Returns:
+            Length of the dataset.
+        """
+        return len(self._dataset)
+
+    def __getitem__(self, index: int):
+        """
+        Returns a single sample.
+
+        Args:
+            index: index of the sample to return.
+        """
+        # Get image and target.
+        img, target = self._dataset.__getitem__(index)
+
+        # Return sample.
+        return index, img, target, self._ix_to_word[target]
+
+    @property
+    def ix_to_word(self):
+        """
+        Returns:
+            Dictionary with mapping of target indices (int) to labels (class names as strings)
+            that can we used by other modules.
+        """
+        return self._ix_to_word

tts/nemo/collections/cv/losses/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.cv.losses.nll_loss import NLLLoss

tts/nemo/collections/cv/losses/nll_loss.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from torch.nn import NLLLoss as torch_NLLLoss
+
+from nemo.core.classes import Serialization, Typing, typecheck
+from nemo.core.neural_types import ClassificationTarget, LogprobsType, LossType, NeuralType
+from nemo.utils.decorators import experimental
+
+
+@experimental
+class NLLLoss(torch_NLLLoss, Serialization, Typing):
+    """ Class representing a simple NLL loss. """
+
+    def __init__(self, name: Optional[str] = None):
+        """
+        Constructor.
+
+        Args:
+            name: Name of the module (DEFAULT: None)
+        """
+        # Call the base constructors.
+        # Serialization.__init__(self, name=name)
+        torch_NLLLoss.__init__(self)
+
+    @property
+    def input_types(self):
+        """ Returns definitions of module input ports. """
+        return {
+            "predictions": NeuralType(axes=('B', 'ANY'), elements_type=LogprobsType()),
+            "targets": NeuralType(axes=('B'), elements_type=ClassificationTarget()),
+        }
+
+    @property
+    def output_types(self):
+        """ Returns definitions of module output ports. """
+        return {"loss": NeuralType(elements_type=LossType())}
+
+    @typecheck()
+    def forward(self, predictions, targets):
+        return torch_NLLLoss().forward(input=predictions, target=targets)

tts/nemo/collections/cv/models/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.cv.models.mnist_lenet5 import MNISTLeNet5, MNISTLeNet5Config

tts/nemo/collections/cv/models/mnist_lenet5.py

@@ -0,0 +1,157 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from torch.utils.data import DataLoader
+
+from nemo.collections.cv.datasets import MNISTDataset, MNISTDatasetConfig
+from nemo.collections.cv.losses import NLLLoss
+from nemo.collections.cv.modules import LeNet5 as LeNet5Module
+from nemo.core.classes import ModelPT
+from nemo.core.classes.common import typecheck
+from nemo.core.config import Config, CosineAnnealingParams, DataLoaderConfig, NovogradParams
+from nemo.core.neural_types import *
+
+
+@dataclass
+class NovogradScheduler(Config):
+    """
+    Scheduler setup for novograd
+    """
+
+    # Mandatory parameters
+    name: str = "CosineAnnealing"
+    args: CosineAnnealingParams = CosineAnnealingParams()  # can be a string to param, or "auto", just like in YAML
+
+    # pytorch lightning parameters
+    monitor: str = "val_loss"
+    iters_per_batch: Optional[float] = None  # computed at runtime
+    max_steps: Optional[int] = None  # computed at runtime or explicitly set here
+    reduce_on_plateau: bool = False
+
+
+@dataclass
+class MNISTOptimizer(Config):
+    """
+    Optimizer setup for novograd
+    """
+
+    # Mandatory arguments
+    name: str = "novograd"
+    lr: float = 0.01
+
+    args: NovogradParams = NovogradParams(betas=(0.8, 0.5))
+    # sched: NovogradScheduler = NovogradScheduler()
+
+
+@dataclass
+class MNISTLeNet5Config(Config):
+    """
+    Structured config for LeNet-5 model class - that also contains parameters of dataset and dataloader.
+    """
+
+    dataset: MNISTDatasetConfig = MNISTDatasetConfig(width=32, height=32)
+    dataloader: DataLoaderConfig = DataLoaderConfig(batch_size=64, shuffle=True)
+    module: Config = Config()
+    optim: MNISTOptimizer = MNISTOptimizer()
+
+
+class MNISTLeNet5(ModelPT):
+    """
+    The LeNet-5 convolutional model.
+    """
+
+    def __init__(self, cfg: MNISTLeNet5Config = MNISTLeNet5Config()):
+        super().__init__(cfg=cfg)
+
+        # Initialize modules.
+        self.module = LeNet5Module(cfg.module)
+        self.loss = NLLLoss()
+
+    @property
+    def input_types(self) -> Optional[Dict[str, NeuralType]]:
+        """
+        Returns:
+            :class:`LeNet5Module` input types.
+        """
+        return self.module.input_types
+
+    @property
+    def output_types(self) -> Optional[Dict[str, NeuralType]]:
+        """
+        Returns:
+            :class:`LeNet5Module` output types.
+        """
+        return self.module.output_types
+
+    @typecheck()
+    def forward(self, images):
+        """ Propagates data by calling the module :class:`LeNet5Module` forward. """
+        return self.module.forward(images=images)
+
+    def setup_training_data(self, train_data_layer_config: Optional[Dict] = None):
+        """ Creates dataset, wrap it with dataloader and return the latter """
+        # Instantiate dataset.
+        mnist_ds = MNISTDataset(self._cfg.dataset)
+        # Configure data loader.
+        train_dataloader = DataLoader(dataset=mnist_ds, **(self._cfg.dataloader))
+        self._train_dl = train_dataloader
+
+    def setup_validation_data(self, val_data_layer_config: Optional[Dict] = None):
+        """ Not implemented. """
+        self._val_dl = None
+
+    def setup_test_data(self, test_data_layer_params: Optional[Dict] = None):
+        """ Not implemented. """
+        self._test_dl = None
+
+    def training_step(self, batch, what_is_this_input):
+        """ Training step, calculate loss. """
+        # "Unpack" the batch.
+        _, images, targets, _ = batch
+
+        # Get predictions.
+        predictions = self(images=images)
+
+        # Calculate loss.
+        loss = self.loss(predictions=predictions, targets=targets)
+
+        # Return it.
+        return {"loss": loss}
+        # of course "return loss" doesn't work :]
+
+    def train_dataloader(self):
+        """ Not implemented. """
+        return self._train_dl
+
+    def save_to(self, save_path: str):
+        """ Not implemented. """
+        pass
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        """ Not implemented. """
+        pass
+
+    @classmethod
+    def list_available_models(cls) -> Optional[Dict[str, str]]:
+        """ Not implemented. """
+        pass
+
+    @classmethod
+    def from_pretrained(cls, name: str):
+        """ Not implemented. """
+        pass

tts/nemo/collections/cv/modules/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo.collections.cv.modules.lenet5 import LeNet5

tts/nemo/collections/cv/modules/lenet5.py

@@ -0,0 +1,114 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Optional
+
+from torch import nn
+
+from nemo.core.classes import NeuralModule, typecheck
+from nemo.core.config import Config
+from nemo.core.neural_types import AxisKind, AxisType, ImageValue, LogprobsType, NeuralType
+
+
+class LeNet5(NeuralModule):
+    """
+    Classical LeNet-5 model for MNIST image classification.
+    """
+
+    def __init__(self, cfg: Config = Config()):
+        """
+        Creates the LeNet-5 model.
+
+        Args:
+            cfg: Default NeMo config containing name.
+        """
+        # Call the base class constructor.
+        super().__init__()  # name=name)
+
+        # Create the LeNet-5 model.
+        self.model = nn.Sequential(
+            nn.Conv2d(1, 6, kernel_size=(5, 5)),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=(2, 2), stride=2),
+            nn.Conv2d(6, 16, kernel_size=(5, 5)),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=(2, 2), stride=2),
+            nn.Conv2d(16, 120, kernel_size=(5, 5)),
+            nn.ReLU(),
+            # reshape to [-1, 120]
+            nn.Flatten(),
+            nn.Linear(120, 84),
+            nn.ReLU(),
+            nn.Linear(84, 10),
+            nn.LogSoftmax(dim=1),
+        )
+
+    @property
+    def input_types(self):
+        """ Returns definitions of module input ports. """
+        return {
+            "images": NeuralType(
+                axes=(
+                    AxisType(kind=AxisKind.Batch),
+                    AxisType(kind=AxisKind.Channel, size=1),
+                    AxisType(kind=AxisKind.Height, size=32),
+                    AxisType(kind=AxisKind.Width, size=32),
+                ),
+                elements_type=ImageValue(),
+            )
+        }
+
+    @property
+    def output_types(self):
+        """ Returns definitions of module output ports. """
+        return {
+            "predictions": NeuralType(
+                axes=(AxisType(kind=AxisKind.Batch), AxisType(kind=AxisKind.Dimension)), elements_type=LogprobsType()
+            )
+        }
+
+    @typecheck()
+    def forward(self, images):
+        """
+        Performs the forward step of the LeNet-5 model.
+
+        Args:
+            images: Batch of images to be classified.
+
+        Returns:
+            Batch of predictions.
+        """
+
+        predictions = self.model(images)
+        return predictions
+
+    def save_to(self, save_path: str):
+        """Not implemented yet.
+           Serialize model.
+
+        Args:
+            save_path (str): path to save serialization.
+        """
+        pass
+
+    @classmethod
+    def restore_from(cls, restore_path: str):
+        """ Not implemented yet.
+            Restore module from serialization.
+
+        Args:
+            restore_path (str): path to serialization
+        """
+        pass