Repository: Lightning-AI/torchmetrics

Issue:
Batch size dependent FID
## 🐛 Bug

As pointed out by @kimihailv in #1198, FID seem to be batch-size dependent.
After some experimentation, the dependency seems to be linked to the inception network.

### To Reproduce

If one runs:

```python
inception = NoTrainInceptionV3(name="inception-v3-compat", features_list=["2048"])
inception.to(device)
imgs = torch.randn(100, 3, 256, 256).to(device)
imgs = ((imgs.clamp(-1, 1) / 2 + 0.5) * 255).to(torch.uint8)
features = []
for img in imgs:
    feature = inception(img.unsqueeze(0))
    features.append(feature)
features_b_1 = torch.cat(features, dim=0)
features = inception(imgs)
```

We then have:

```python
torch.allclose(features, features_b_1)
>> False
torch.norm(features - features_b_1, p=2)
>> tensor(0.2950, device='cuda:0')
```

To be noted that when replacing 
```python
features = inception(imgs)
```

by 
```python
features = []
for img in imgs:
    feature = inception(img.unsqueeze(0))
    features.append(feature)
features= torch.cat(features, dim=0)
```
the FID is again batch independent. However, this is not a possible fix. Indeed, first this is very ineficient. Also, from experimentation, the batch bias in FID seems to be higher from small batch-sizes. If we compute FID between 2 uniformly sampled distributions with a 1000 points each, if we compute it with a batch size of 1000, we get FID 1.9 but if we compute it with batch-size=2, then the FID is 10. Since we sample from the same distribution, FID should be as close too zero as possible.

### Expected behavior

FID computation should be batch_size independent

### Environment

- TorchMetrics version (and how you installed TM, e.g. `conda`, `pip`, build from source):  0.10.3
- Python & PyTorch Version (e.g., 1.0): 3.10 and 1.13
- Any other relevant information such as OS (e.g., Linux): Linux



Discussion:
Hi @nicolas-dufour, thanks for reporting this issue.
I tried reproducing your results, but were unable to do that, please see this notebook:
https://colab.research.google.com/drive/16aVuIJQj7TUmhiy6TGL2QcyQOTb3ibJr?usp=sharing
On CPU i get a norm difference of `tensor(0.0002)` and on CUDA i get `tensor(0.0172)` between the two approaches, which seems reasonable to still calculate the metric correctly.
Hi @SkafteNicki, thanks for checking this out.

Hum that is strange! From further experiments, i found that the discrepancy disappeared when using float64. Maybe the problem is accelerator dependent? The previous experiment was done on a RTX 3090. 

Also, I've observed that the impact on FID was minimal for dataset size > 100. However, it's still weird that the metric changes with respect to the batch size. One solution would be to offer the option to run the embedding network at float64 precision.

Resolving patch:
diff --git a/src/torchmetrics/image/fid.py b/src/torchmetrics/image/fid.py
--- a/src/torchmetrics/image/fid.py
+++ b/src/torchmetrics/image/fid.py
@@ -19,6 +19,7 @@
 from torch import Tensor
 from torch.autograd import Function
 from torch.nn import Module
+from torch.nn.functional import adaptive_avg_pool2d
 
 from torchmetrics.metric import Metric
 from torchmetrics.utilities import rank_zero_info
@@ -30,11 +31,16 @@
 
 if _TORCH_FIDELITY_AVAILABLE:
     from torch_fidelity.feature_extractor_inceptionv3 import FeatureExtractorInceptionV3 as _FeatureExtractorInceptionV3
+    from torch_fidelity.helpers import vassert
+    from torch_fidelity.interpolate_compat_tensorflow import interpolate_bilinear_2d_like_tensorflow1x
 else:
 
     class _FeatureExtractorInceptionV3(Module):
         pass
 
+    vassert = None
+    interpolate_bilinear_2d_like_tensorflow1x = None
+
     __doctest_skip__ = ["FrechetInceptionDistance", "FrechetInceptionDistance.plot"]
 
 
@@ -59,9 +65,94 @@ def train(self, mode: bool) -> "NoTrainInceptionV3":
         """Force network to always be in evaluation mode."""
         return super().train(False)
 
+    def _torch_fidelity_forward(self, x: Tensor) -> Tensor:
+        """Forward method of inception net.
+
+        Copy of the forward method from this file:
+        https://github.com/toshas/torch-fidelity/blob/master/torch_fidelity/feature_extractor_inceptionv3.py
+        with a single line change regarding the casting of `x` in the beginning.
+
+        Corresponding license file (Apache License, Version 2.0):
+        https://github.com/toshas/torch-fidelity/blob/master/LICENSE.md
+        """
+        vassert(torch.is_tensor(x) and x.dtype == torch.uint8, "Expecting image as torch.Tensor with dtype=torch.uint8")
+        features = {}
+        remaining_features = self.features_list.copy()
+
+        x = x.to(self._dtype) if hasattr(self, "_dtype") else x.to(torch.float)
+        x = interpolate_bilinear_2d_like_tensorflow1x(
+            x,
+            size=(self.INPUT_IMAGE_SIZE, self.INPUT_IMAGE_SIZE),
+            align_corners=False,
+        )
+        x = (x - 128) / 128
+
+        x = self.Conv2d_1a_3x3(x)
+        x = self.Conv2d_2a_3x3(x)
+        x = self.Conv2d_2b_3x3(x)
+        x = self.MaxPool_1(x)
+
+        if "64" in remaining_features:
+            features["64"] = adaptive_avg_pool2d(x, output_size=(1, 1)).squeeze(-1).squeeze(-1)
+            remaining_features.remove("64")
+            if len(remaining_features) == 0:
+                return tuple(features[a] for a in self.features_list)
+
+        x = self.Conv2d_3b_1x1(x)
+        x = self.Conv2d_4a_3x3(x)
+        x = self.MaxPool_2(x)
+
+        if "192" in remaining_features:
+            features["192"] = adaptive_avg_pool2d(x, output_size=(1, 1)).squeeze(-1).squeeze(-1)
+            remaining_features.remove("192")
+            if len(remaining_features) == 0:
+                return tuple(features[a] for a in self.features_list)
+
+        x = self.Mixed_5b(x)
+        x = self.Mixed_5c(x)
+        x = self.Mixed_5d(x)
+        x = self.Mixed_6a(x)
+        x = self.Mixed_6b(x)
+        x = self.Mixed_6c(x)
+        x = self.Mixed_6d(x)
+        x = self.Mixed_6e(x)
+
+        if "768" in remaining_features:
+            features["768"] = adaptive_avg_pool2d(x, output_size=(1, 1)).squeeze(-1).squeeze(-1)
+            remaining_features.remove("768")
+            if len(remaining_features) == 0:
+                return tuple(features[a] for a in self.features_list)
+
+        x = self.Mixed_7a(x)
+        x = self.Mixed_7b(x)
+        x = self.Mixed_7c(x)
+        x = self.AvgPool(x)
+        x = torch.flatten(x, 1)
+
+        if "2048" in remaining_features:
+            features["2048"] = x
+            remaining_features.remove("2048")
+            if len(remaining_features) == 0:
+                return tuple(features[a] for a in self.features_list)
+
+        if "logits_unbiased" in remaining_features:
+            x = x.mm(self.fc.weight.T)
+            # N x 1008 (num_classes)
+            features["logits_unbiased"] = x
+            remaining_features.remove("logits_unbiased")
+            if len(remaining_features) == 0:
+                return tuple(features[a] for a in self.features_list)
+
+            x = x + self.fc.bias.unsqueeze(0)
+        else:
+            x = self.fc(x)
+
+        features["logits"] = x
+        return tuple(features[a] for a in self.features_list)
+
     def forward(self, x: Tensor) -> Tensor:
         """Forward pass of neural network with reshaping of output."""
-        out = super().forward(x)
+        out = self._torch_fidelity_forward(x)
         return out[0].reshape(x.shape[0], -1)
 
 
@@ -151,6 +242,10 @@ class FrechetInceptionDistance(Metric):
     flag ``real`` determines if the images should update the statistics of the real distribution or the
     fake distribution.
 
+    This metric is known to be unstable in its calculatations, and we recommend for the best results using this metric
+    that you calculate using `torch.float64` (default is `torch.float32`) which can be set using the `.set_dtype`
+    method of the metric.
+
     .. note:: using this metrics requires you to have ``scipy`` install. Either install as ``pip install
         torchmetrics[image]`` or ``pip install scipy``
 
@@ -307,6 +402,17 @@ def reset(self) -> None:
         else:
             super().reset()
 
+    def set_dtype(self, dst_type: Union[str, torch.dtype]) -> "Metric":
+        """Transfer all metric state to specific dtype. Special version of standard `type` method.
+
+        Arguments:
+            dst_type (type or string): the desired type.
+        """
+        out = super().set_dtype(dst_type)
+        if isinstance(out.inception, NoTrainInceptionV3):
+            out.inception._dtype = dst_type
+        return out
+
     def plot(
         self, val: Optional[Union[Tensor, Sequence[Tensor]]] = None, ax: Optional[_AX_TYPE] = None
     ) -> _PLOT_OUT_TYPE:


Test patch:
diff --git a/tests/unittests/image/test_fid.py b/tests/unittests/image/test_fid.py
--- a/tests/unittests/image/test_fid.py
+++ b/tests/unittests/image/test_fid.py
@@ -199,3 +199,15 @@ def test_normalize_arg_false():
     metric = FrechetInceptionDistance(normalize=False)
     with pytest.raises(ValueError, match="Expecting image as torch.Tensor with dtype=torch.uint8"):
         metric.update(img, real=True)
+
+
+def test_dtype_transfer_to_submodule():
+    """Test that change in dtype also changes the default inception net."""
+    imgs = torch.randn(1, 3, 256, 256)
+    imgs = ((imgs.clamp(-1, 1) / 2 + 0.5) * 255).to(torch.uint8)
+
+    metric = FrechetInceptionDistance(feature=64)
+    metric.set_dtype(torch.float64)
+
+    out = metric.inception(imgs)
+    assert out.dtype == torch.float64