Introduce code from Rust monorepo. This is a mess

tts/nemo/collections/common/metrics/__init__.py

@@ -0,0 +1,17 @@
+# 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.common.metrics.classification_accuracy import TopKClassificationAccuracy
+from nemo.collections.common.metrics.global_average_loss_metric import GlobalAverageLossMetric
+from nemo.collections.common.metrics.perplexity import Perplexity

tts/nemo/collections/common/metrics/classification_accuracy.py

@@ -0,0 +1,152 @@
+# 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 List
+
+import torch
+from torchmetrics import Metric
+
+__all__ = ['TopKClassificationAccuracy']
+
+
+class TopKClassificationAccuracy(Metric):
+    """
+    This metric computes numerator and denominator for Overall Accuracy between logits and labels.
+    When doing distributed training/evaluation the result of res=TopKClassificationAccuracy(logits, labels) calls
+    will be all-reduced between all workers using SUM operations.
+    Here contains two numbers res=[correctly_predicted, total_samples]. Accuracy=correctly_predicted/total_samples.
+
+    If used with PytorchLightning LightningModule, include correct_count and total_count inside validation_step results.
+    Then aggregate (sum) then at the end of validation epoch to correctly compute validation WER.
+
+    Example:
+        def validation_step(self, batch, batch_idx):
+            ...
+            correct_count, total_count = self._accuracy(logits, labels)
+            return {'val_loss': loss_value, 'val_correct_count': correct_count, 'val_total_count': total_count}
+
+        def validation_epoch_end(self, outputs):
+            ...
+            val_loss_mean = torch.stack([x['val_loss'] for x in outputs]).mean()
+            correct_counts = torch.stack([x['val_correct_counts'] for x in outputs])
+            total_counts = torch.stack([x['val_total_counts'] for x in outputs])
+
+            topk_scores = compute_topk_accuracy(correct_counts, total_counts)
+
+            tensorboard_log = {'val_loss': val_loss_mean}
+            for top_k, score in zip(self._accuracy.top_k, topk_scores):
+                tensorboard_log['val_epoch_top@{}'.format(top_k)] = score
+
+            return {'log': tensorboard_log}
+
+    Args:
+        top_k: Optional list of integers. Defaults to [1].
+
+    Returns:
+        res: a torch.Tensor object with two elements: [correct_count, total_count]. To correctly compute average
+        accuracy, compute acc=correct_count/total_count
+    """
+
+    def __init__(self, top_k=None, dist_sync_on_step=False):
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+
+        if top_k is None:
+            top_k = [1]
+
+        self.top_k = top_k
+        self.add_state(
+            "correct_counts_k", default=torch.zeros(len(self.top_k)), dist_reduce_fx='sum', persistent=False
+        )
+        self.add_state("total_counts_k", default=torch.zeros(len(self.top_k)), dist_reduce_fx='sum', persistent=False)
+
+    @torch.no_grad()
+    def top_k_predicted_labels(self, logits: torch.Tensor) -> torch.Tensor:
+        max_k = max(self.top_k)
+        _, predictions = logits.topk(max_k, dim=1, largest=True, sorted=True)
+        return predictions
+
+    def update(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        with torch.no_grad():
+            predictions = self.top_k_predicted_labels(logits)
+            predictions = predictions.t()
+            correct = predictions.eq(labels.view(1, -1)).expand_as(predictions)
+
+            correct_counts_k = []
+            total_counts_k = []
+
+            for k in self.top_k:
+                correct_k = correct[:k].reshape(-1).long().sum()
+                total_k = labels.shape[0]
+
+                correct_counts_k.append(correct_k)
+                total_counts_k.append(total_k)
+
+            self.correct_counts_k = torch.tensor(correct_counts_k, dtype=labels.dtype, device=labels.device)
+            self.total_counts_k = torch.tensor(total_counts_k, dtype=labels.dtype, device=labels.device)
+
+    def compute(self):
+        """
+        Computes the top-k accuracy.
+
+        Returns:
+            A list of length `K`, such that k-th index corresponds to top-k accuracy
+            over all distributed processes.
+        """
+        if not len(self.correct_counts_k) == len(self.top_k) == len(self.total_counts_k):
+            raise ValueError("length of counts must match to topk length")
+
+        if self.top_k == [1]:
+            return [self.correct_counts_k.float() / self.total_counts_k]
+
+        else:
+            top_k_scores = compute_topk_accuracy(self.correct_counts_k, self.total_counts_k)
+
+            return top_k_scores
+
+    @property
+    def top_k(self) -> List[int]:
+        return self._top_k
+
+    @top_k.setter
+    def top_k(self, value: List[int]):
+        if value is None:
+            value = [1]
+
+        if type(value) == int:
+            value = [value]
+
+        if type(value) != list:
+            value = list(value)
+
+        self._top_k = value
+
+
+def compute_topk_accuracy(correct_counts_k, total_counts_k):
+    """
+    Computes the top-k accuracy
+    Args:
+        correct_counts: Tensor of shape [K], K being the top-k parameter.
+        total_counts: Tensor of shape [K], and K being the top-k parameter.
+    Returns:
+        A list of length `K`, such that k-th index corresponds to top-k accuracy
+        over all distributed processes.
+    """
+    top_k_scores = []
+
+    for ki in range(len(correct_counts_k)):
+        correct_count = correct_counts_k[ki].item()
+        total_count = total_counts_k[ki].item()
+        top_k_scores.append(correct_count / float(total_count))
+
+    return top_k_scores

tts/nemo/collections/common/metrics/global_average_loss_metric.py

@@ -0,0 +1,73 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torchmetrics import Metric
+
+__all__ = ['GlobalAverageLossMetric']
+
+
+class GlobalAverageLossMetric(Metric):
+    """
+    This class is for averaging loss across multiple processes if a distributed backend is used. True average is
+    computed not running average. It does not accumulate gradients so the averaged loss cannot be used for optimization.
+    If ``take_avg_loss`` is ``True``, the :meth:`update` method ``loss`` argument has to be a mean loss. If
+    ``take_avg_loss`` is ``False`` then the :meth:`update` method ``loss`` argument has to be a sum of losses.
+
+    See :doc:`PyTorch Lightning Metrics<pytorch-lightning:metrics>` for the metric usage instruction.
+
+    Args:
+        compute_on_step:
+            Forward only calls ``update()`` and returns ``None`` if this is set to ``False``. default: ``True``
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step.
+        process_group:
+            Specify the process group on which synchronization is called. default: ``None`` (which selects the entire
+                world)
+        take_avg_loss:
+            If ``True`` values of :meth:`update` method ``loss`` argument has to be a mean loss. If ``False``
+            values of :meth:`update` method ``loss`` argument has to be a sum of losses. default: ``True``
+    """
+
+    def __init__(self, compute_on_step=True, dist_sync_on_step=False, process_group=None, take_avg_loss=True):
+        super().__init__(
+            compute_on_step=compute_on_step, dist_sync_on_step=dist_sync_on_step, process_group=process_group
+        )
+        self.add_state("loss_sum", torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx='sum')
+        self.add_state("num_measurements", torch.tensor(0, dtype=torch.int64), dist_reduce_fx='sum')
+        self.take_avg_loss = take_avg_loss
+
+    def update(self, loss, num_measurements):
+        """
+        Updates :attr:`loss_sum` and :attr:`num_measurements`.
+
+        Args:
+            loss: A float zero dimensional ``torch.Tensor`` which is either sum or average of losses for processed
+                examples. See ``take_avg_loss`` parameter of :meth:`__init__`.
+            num_measurements: An integer zero dimensional ``torch.Tensor`` which contains a number of loss measurements.
+                The sum or mean of the results of these measurements are in the ``loss`` parameter.
+        """
+        if self.take_avg_loss:
+            loss *= num_measurements
+        self.loss_sum += loss
+        self.num_measurements += num_measurements
+
+    def compute(self):
+        """
+        Returns mean loss.
+        """
+        if self.num_measurements.eq(0):
+            return torch.tensor(float('nan'))
+        return self.loss_sum / self.num_measurements

tts/nemo/collections/common/metrics/perplexity.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch.distributions.categorical import Categorical
+from torchmetrics import Metric
+
+__all__ = ['Perplexity']
+
+
+class Perplexity(Metric):
+    """
+    This class computes mean perplexity of distributions in the last dimension of inputs. It is a wrapper around
+    :doc:`torch.distributions.Categorical.perplexity<pytorch:distributions>` method. You have to provide either 
+    ``probs`` or ``logits`` to the :meth:`update` method. The class computes perplexities for distributions passed to 
+    :meth:`update` method in ``probs`` or ``logits`` arguments and averages the perplexities. Reducing results between
+    all workers is done via SUM operations.
+    See :doc:`PyTorch Lightning Metrics<pytorch-lightning:metrics>` for the metric usage instructions.
+    Args:
+        compute_on_step:
+            Forward only calls ``update()`` and returns ``None`` if this is set to ``False``. default: ``True``
+        dist_sync_on_step:
+            Synchronize metric state across processes at each ``forward()``
+            before returning the value at the step.
+        process_group:
+            Specify the process group on which synchronization is called. default: ``None`` (which selects the entire
+                world)
+        validate_args:
+            If ``True`` values of :meth:`update` method parameters are checked. ``logits`` has to not contain NaNs and
+            ``probs`` last dim has to be valid probability distribution.
+    """
+
+    def __init__(self, compute_on_step=True, dist_sync_on_step=False, process_group=None, validate_args=True):
+        super().__init__(
+            compute_on_step=compute_on_step, dist_sync_on_step=dist_sync_on_step, process_group=process_group
+        )
+        self.validate_args = validate_args
+        self.add_state('perplexities_sum', torch.tensor(0.0, dtype=torch.float64), dist_reduce_fx='sum')
+        # Total number of distributions seen since last reset
+        self.add_state('num_distributions', torch.tensor(0, dtype=torch.int64), dist_reduce_fx='sum')
+
+    def update(self, probs=None, logits=None):
+        """
+        Updates :attr:`perplexities_sum` and :attr:`num_distributions`.
+        Args:
+            probs: A ``torch.Tensor`` which innermost dimension is valid probability distribution.
+            logits: A ``torch.Tensor`` without NaNs.
+        """
+        d = Categorical(probs, logits, validate_args=self.validate_args)
+        ppl = d.perplexity()
+        self.num_distributions += ppl.numel()
+        self.perplexities_sum += ppl.sum()
+
+    def compute(self):
+        """
+        Returns perplexity across all workers and resets to 0 :attr:`perplexities_sum` and :attr:`num_distributions`.
+        """
+        if self.num_distributions.eq(0):
+            return None
+        return self.perplexities_sum / self.num_distributions