Inception Score

docs/source/metrics.rst

@@ -334,6 +334,7 @@ Complete list of metrics
     Rouge
     RougeL
     RougeN
+    InceptionScore
     FID
 
 Helpers for customizing metrics

ignite/metrics/__init__.py

@@ -6,6 +6,7 @@
 from ignite.metrics.fbeta import Fbeta
 from ignite.metrics.frequency import Frequency
 from ignite.metrics.gan.fid import FID
+from ignite.metrics.gan.inception_score import InceptionScore
 from ignite.metrics.loss import Loss
 from ignite.metrics.mean_absolute_error import MeanAbsoluteError
 from ignite.metrics.mean_pairwise_distance import MeanPairwiseDistance
@@ -41,6 +42,7 @@
     "FID",
     "GeometricAverage",
     "IoU",
+    "InceptionScore",
     "mIoU",
     "JaccardIndex",
     "MultiLabelConfusionMatrix",

ignite/metrics/gan/__init__.py

@@ -1,5 +1,7 @@
 from ignite.metrics.gan.fid import FID
+from ignite.metrics.gan.inception_score import InceptionScore
 
 __all__ = [
+    "InceptionScore",
     "FID",
 ]

ignite/metrics/gan/fid.py

@@ -4,7 +4,8 @@
 
 import torch
 
-from ignite.metrics.metric import Metric, reinit__is_reduced, sync_all_reduce
+from ignite.metrics.gan.utils import InceptionModel, _BaseInceptionMetric
+from ignite.metrics.metric import reinit__is_reduced, sync_all_reduce
 
 __all__ = [
     "FID",
@@ -47,26 +48,7 @@ def fid_score(
     return float(diff.dot(diff).item() + torch.trace(sigma1) + torch.trace(sigma2) - 2 * tr_covmean)
 
 
-class InceptionExtractor:
-    def __init__(self) -> None:
-        try:
-            from torchvision import models
-        except ImportError:
-            raise RuntimeError("This module requires torchvision to be installed.")
-        self.model = models.inception_v3(pretrained=True)
-        self.model.fc = torch.nn.Identity()
-        self.model.eval()
-
-    @torch.no_grad()
-    def __call__(self, data: torch.Tensor) -> torch.Tensor:
-        if data.dim() != 4:
-            raise ValueError(f"Inputs should be a tensor of dim 4, got {data.dim()}")
-        if data.shape[1] != 3:
-            raise ValueError(f"Inputs should be a tensor with 3 channels, got {data.shape}")
-        return self.model(data)
-
-
-class FID(Metric):
+class FID(_BaseInceptionMetric):
     r"""Calculates Frechet Inception Distance.
 
     .. math::
@@ -90,10 +72,15 @@ class FID(Metric):
         __ https://github.com/mseitzer/pytorch-fid
 
     Args:
-        num_features: number of features, must be defined if the parameter ``feature_extractor`` is also defined.
-            Otherwise, default value is 2048.
-        feature_extractor: a callable for extracting the features from the input data. If neither num_features nor
-            feature_extractor are defined, default value is ``InceptionExtractor``.
+        num_features: number of features predicted by the model or the reduced feature vector of the image.
+            Default value is 2048.
+        feature_extractor: a torch Module for extracting the features from the input data.
+            It returns a tensor of shape (batch_size, num_features).
+            If neither ``num_features`` nor ``feature_extractor`` are defined, by default we use an ImageNet
+            pretrained Inception Model. If only ``num_features`` is defined but ``feature_extractor`` is not
+            defined, ``feature_extractor`` is assigned Identity Function.
+            Please note that the model will be implicitly converted to device mentioned in the ``device``
+            argument.
         output_transform: a callable that is used to transform the
             :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the
             form expected by the metric. This can be useful if, for example, you have a multi-output model and
@@ -110,7 +97,7 @@ class FID(Metric):
             import torch
             from ignite.metric.gan import FID
 
-            y_pred, y = torch.rand(10, 2048), torch.rand(10, 2048)
+            y_pred, y = torch.rand(10, 3, 299, 299), torch.rand(10, 3, 299, 299)
             m = FID()
             m.update((y_pred, y))
             print(m.compute())
@@ -121,7 +108,7 @@ class FID(Metric):
     def __init__(
         self,
         num_features: Optional[int] = None,
-        feature_extractor: Optional[Callable] = None,
+        feature_extractor: Optional[torch.nn.Module] = None,
         output_transform: Callable = lambda x: x,
         device: Union[str, torch.device] = torch.device("cpu"),
     ) -> None:
@@ -136,26 +123,24 @@ def __init__(
         except ImportError:
             raise RuntimeError("This module requires scipy to be installed.")
 
-        # default is inception
         if num_features is None and feature_extractor is None:
-            num_features = 2048
-            feature_extractor = InceptionExtractor()
-        elif num_features is None:
-            raise ValueError("Argument num_features should be defined")
-        elif feature_extractor is None:
-            self._feature_extractor = lambda x: x
-            feature_extractor = self._feature_extractor
-
-        if num_features <= 0:
-            raise ValueError(f"Argument num_features must be greater to zero, got: {num_features}")
-        self._num_features = num_features
-        self._feature_extractor = feature_extractor
+            num_features = 1000
+            feature_extractor = InceptionModel(return_features=False)
+
         self._eps = 1e-6
-        super(FID, self).__init__(output_transform=output_transform, device=device)
+
+        super(FID, self).__init__(
+            num_features=num_features,
+            feature_extractor=feature_extractor,
+            output_transform=output_transform,
+            device=device,
+        )
 
     @staticmethod
     def _online_update(features: torch.Tensor, total: torch.Tensor, sigma: torch.Tensor) -> None:
+
         total += features
+
         if LooseVersion(torch.__version__) <= LooseVersion("1.7.0"):
             sigma += torch.ger(features, features)
         else:
@@ -165,65 +150,70 @@ def _get_covariance(self, sigma: torch.Tensor, total: torch.Tensor) -> torch.Ten
         r"""
         Calculates covariance from mean and sum of products of variables
         """
+
         if LooseVersion(torch.__version__) <= LooseVersion("1.7.0"):
             sub_matrix = torch.ger(total, total)
         else:
             sub_matrix = torch.outer(total, total)
+
         sub_matrix = sub_matrix / self._num_examples
-        return (sigma - sub_matrix) / (self._num_examples - 1)
 
-    @staticmethod
-    def _check_feature_input(train: torch.Tensor, test: torch.Tensor) -> None:
-        for feature in [train, test]:
-            if feature.dim() != 2:
-                raise ValueError(f"Features must be a tensor of dim 2, got: {feature.dim()}")
-            if feature.shape[0] == 0:
-                raise ValueError(f"Batch size should be greater than one, got: {feature.shape[0]}")
-            if feature.shape[1] == 0:
-                raise ValueError(f"Feature size should be greater than one, got: {feature.shape[1]}")
-        if train.shape[0] != test.shape[0] or train.shape[1] != test.shape[1]:
-            raise ValueError(
-                f"Number of Training Features and Testing Features should be equal ({train.shape} != {test.shape})"
-            )
+        return (sigma - sub_matrix) / (self._num_examples - 1)
 
     @reinit__is_reduced
     def reset(self) -> None:
-        self._train_sigma = torch.zeros((self._num_features, self._num_features), dtype=torch.float64).to(self._device)
-        self._train_total = torch.zeros(self._num_features, dtype=torch.float64).to(self._device)
-        self._test_sigma = torch.zeros((self._num_features, self._num_features), dtype=torch.float64).to(self._device)
-        self._test_total = torch.zeros(self._num_features, dtype=torch.float64).to(self._device)
-        self._num_examples = 0
+
+        self._train_sigma = torch.zeros(
+            (self._num_features, self._num_features), dtype=torch.float64, device=self._device
+        )
+
+        self._train_total = torch.zeros(self._num_features, dtype=torch.float64, device=self._device)
+
+        self._test_sigma = torch.zeros(
+            (self._num_features, self._num_features), dtype=torch.float64, device=self._device
+        )
+
+        self._test_total = torch.zeros(self._num_features, dtype=torch.float64, device=self._device)
+        self._num_examples: int = 0
+
         super(FID, self).reset()
 
     @reinit__is_reduced
     def update(self, output: Sequence[torch.Tensor]) -> None:
 
-        # Extract the features from the outputs
-        train_features = self._feature_extractor(output[0].detach()).to(self._device)
-        test_features = self._feature_extractor(output[1].detach()).to(self._device)
+        train, test = output
+        train_features = self._extract_features(train)
+        test_features = self._extract_features(test)
 
-        # Check the feature shapess
-        self._check_feature_input(train_features, test_features)
+        if train_features.shape[0] != test_features.shape[0] or train_features.shape[1] != test_features.shape[1]:
+            raise ValueError(
+                f"""
+    Number of Training Features and Testing Features should be equal ({train_features.shape} != {test_features.shape})
+                """
+            )
 
         # Updates the mean and covariance for the train features
-        for i, features in enumerate(train_features, start=self._num_examples + 1):
+        for features in train_features:
             self._online_update(features, self._train_total, self._train_sigma)
 
         # Updates the mean and covariance for the test features
-        for i, features in enumerate(test_features, start=self._num_examples + 1):
+        for features in test_features:
             self._online_update(features, self._test_total, self._test_sigma)
 
         self._num_examples += train_features.shape[0]
 
     @sync_all_reduce("_num_examples", "_train_total", "_test_total", "_train_sigma", "_test_sigma")
     def compute(self) -> float:
+
         fid = fid_score(
             mu1=self._train_total / self._num_examples,
             mu2=self._test_total / self._num_examples,
             sigma1=self._get_covariance(self._train_sigma, self._train_total),
             sigma2=self._get_covariance(self._test_sigma, self._test_total),
             eps=self._eps,
         )
+
         if torch.isnan(torch.tensor(fid)) or torch.isinf(torch.tensor(fid)):
             warnings.warn("The product of covariance of train and test features is out of bounds.")
+
         return fid

ignite/metrics/gan/inception_score.py

@@ -0,0 +1,115 @@
+from typing import Callable, Optional, Union
+
+import torch
+
+from ignite.exceptions import NotComputableError
+from ignite.metrics.gan.utils import InceptionModel, _BaseInceptionMetric
+
+# These decorators helps with distributed settings
+from ignite.metrics.metric import reinit__is_reduced, sync_all_reduce
+
+__all__ = ["InceptionScore"]
+
+
+class InceptionScore(_BaseInceptionMetric):
+    r"""Calculates Inception Score.
+
+    .. math::
+       \text{IS(G)} = \exp(\frac{1}{N}\sum_{i=1}^{N} D_{KL} (p(y|x^{(i)} \parallel \hat{p}(y))))
+
+    where :math:`p(y|x)` is the conditional probability of image being the given object and
+    :math:`p(y)` is the marginal probability that the given image is real, `G` refers to the
+    generated image and :math:`D_{KL}` refers to KL Divergence of the above mentioned probabilities.
+
+    More details can be found in `Barratt et al. 2018`__.
+
+    __ https://arxiv.org/pdf/1801.01973.pdf
+
+
+    Args:
+        num_features: number of features predicted by the model or number of classes of the model. Default
+            value is 1000.
+        feature_extractor: a torch Module for predicting the probabilities from the input data.
+            It returns a tensor of shape (batch_size, num_features).
+            If neither ``num_features`` nor ``feature_extractor`` are defined, by default we use an ImageNet
+            pretrained Inception Model. If only ``num_features`` is defined but ``feature_extractor`` is not
+            defined, ``feature_extractor`` is assigned Identity Function.
+            Please note that the class object will be implicitly converted to device mentioned in the
+            ``device`` argument.
+        output_transform: a callable that is used to transform the
+            :class:`~ignite.engine.engine.Engine`'s ``process_function``'s output into the
+            form expected by the metric. This can be useful if, for example, you have a multi-output model and
+            you want to compute the metric with respect to one of the outputs.
+            By default, metrics require the output as ``y_pred``.
+        device: specifies which device updates are accumulated on. Setting the
+            metric's device to be the same as your ``update`` arguments ensures the ``update`` method is
+            non-blocking. By default, CPU.
+
+    Example:
+
+        .. code-block:: python
+
+            from ignite.metric.gan import InceptionScore
+            import torch
+
+            images = torch.rand(10, 3, 299, 299)
+
+            m = InceptionScore()
+            m.update(images)
+            print(m.compute())
+
+    .. versionadded:: 0.5.0
+    """
+
+    def __init__(
+        self,
+        num_features: Optional[int] = None,
+        feature_extractor: Optional[torch.nn.Module] = None,
+        output_transform: Callable = lambda x: x,
+        device: Union[str, torch.device] = torch.device("cpu"),
+    ) -> None:
+
+        if num_features is None and feature_extractor is None:
+            num_features = 1000
+            feature_extractor = InceptionModel(return_features=False)
+
+        self._eps = 1e-16
+
+        super(InceptionScore, self).__init__(
+            num_features=num_features,
+            feature_extractor=feature_extractor,
+            output_transform=output_transform,
+            device=device,
+        )
+
+    @reinit__is_reduced
+    def reset(self) -> None:
+
+        self._num_examples = 0
+
+        self._prob_total = torch.zeros(self._num_features, dtype=torch.float64, device=self._device)
+        self._total_kl_d = torch.zeros(self._num_features, dtype=torch.float64, device=self._device)
+
+        super(InceptionScore, self).reset()
+
+    @reinit__is_reduced
+    def update(self, output: torch.Tensor) -> None:
+
+        probabilities = self._extract_features(output)
+
+        self._num_examples += probabilities.shape[0]
+
+        self._prob_total += torch.sum(probabilities, 0).to(self._device)
+        self._total_kl_d += torch.sum(probabilities * torch.log(probabilities + self._eps), 0).to(self._device)
+
+    @sync_all_reduce("_num_examples", "_prob_total", "_total_kl_d")
+    def compute(self) -> torch.Tensor:
+
+        if self._num_examples == 0:
+            raise NotComputableError("InceptionScore must have at least one example before it can be computed.")
+
+        mean_probs = self._prob_total / self._num_examples
+        excess_entropy = self._prob_total * torch.log(mean_probs + self._eps)
+        avg_kl_d = torch.sum(self._total_kl_d - excess_entropy) / self._num_examples
+
+        return torch.exp(avg_kl_d)