Repository: Lightning-AI/torchmetrics

Issue:
FrechetInceptionDistance wrong behavior for calculating two features with different batch size
## 🐛 Bug

I want to calculate FID between the real and fake images, but the batch_size of fake images is 5 times the real ones (according to the traditional setting in layout2img), and the result is wrong. Then I checked the implementation of [FID](https://github.com/PyTorchLightning/metrics/blob/a971c6b456e40728b34494ff9186af20da46cb5b/torchmetrics/image/fid.py#L270), which assumes the same batch_size when calculating the covariance, which is not suitable to my knowledge.

### To Reproduce

``` python
def test_step(self, batch, batch_idx):
        real_imgs, label, bbox, id = batch
        self.FID.update(unit2uint8(real_imgs), real=True)
        fake_list = []
        for _ in range(self.hparams.test_sample):# 5
            z_obj = torch.randn(real_imgs.shape[0], self.hparams.num_obj, self.hparams.z_dim, device=self.device)
            z_im = torch.randn(real_imgs.shape[0], self.hparams.z_dim, device=self.device)
            fake_imgs = self(z_obj, bbox, z_im=z_im, y=label.squeeze(dim=-1))
            fake_list.append(fake_imgs)
        fake_imgs = torch.cat(fake_list, dim=0)
        self.FID.update(unit2uint8(fake_imgs), real=False)
        self.IS.update(unit2uint8(fake_imgs))
```
<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

#### Code sample
Here is the official [implementation](https://github.com/PyTorchLightning/metrics/blob/a971c6b456e40728b34494ff9186af20da46cb5b/torchmetrics/image/fid.py#L260)
```python
    def compute(self) -> Tensor:
        """Calculate FID score based on accumulated extracted features from the two distributions."""
        real_features = dim_zero_cat(self.real_features)
        fake_features = dim_zero_cat(self.fake_features)
        # computation is extremely sensitive so it needs to happen in double precision
        orig_dtype = real_features.dtype
        real_features = real_features.double()
        fake_features = fake_features.double()

        # calculate mean and covariance
        n = real_features.shape[0]
        mean1 = real_features.mean(dim=0)
        mean2 = fake_features.mean(dim=0)
        diff1 = real_features - mean1
        diff2 = fake_features - mean2
        cov1 = 1.0 / (n - 1) * diff1.t().mm(diff1) # the same n is used
        cov2 = 1.0 / (n - 1) * diff2.t().mm(diff2)

        # compute fid
        return _compute_fid(mean1, cov1, mean2, cov2).to(orig_dtype)
```
<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment

- TorchMetrics version ( `pip`): 0.8.0
- PyTorch Version  :1.12.0a0+bd13bc6
- Any other relevant information such as OS (e.g., Linux): docker environment, the image is highly based on nvcr.io/nvidia/pytorch:22.04-py3 

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!
FYI, the more suitable behavior might be something like the following code
```python
    def compute(self) -> Tensor:
        """Calculate FID score based on accumulated extracted features from the two distributions."""
        real_features = dim_zero_cat(self.real_features)
        fake_features = dim_zero_cat(self.fake_features)
        # computation is extremely sensitive so it needs to happen in double precision
        orig_dtype = real_features.dtype
        real_features = real_features.double()
        fake_features = fake_features.double()

        # calculate mean and covariance
        n = real_features.shape[0]
        m = fake_features.shape[0]                  ### new
        mean1 = real_features.mean(dim=0)
        mean2 = fake_features.mean(dim=0)
        diff1 = real_features - mean1
        diff2 = fake_features - mean2
        cov1 = 1.0 / (n - 1) * diff1.t().mm(diff1)  
        cov2 = 1.0 / (m - 1) * diff2.t().mm(diff2)            ### change n to m

        # compute fid
        return _compute_fid(mean1, cov1, mean2, cov2).to(orig_dtype)
```

Resolving patch:
diff --git a/torchmetrics/image/fid.py b/torchmetrics/image/fid.py
--- a/torchmetrics/image/fid.py
+++ b/torchmetrics/image/fid.py
@@ -268,12 +268,13 @@ def compute(self) -> Tensor:
 
         # calculate mean and covariance
         n = real_features.shape[0]
+        m = fake_features.shape[0]
         mean1 = real_features.mean(dim=0)
         mean2 = fake_features.mean(dim=0)
         diff1 = real_features - mean1
         diff2 = fake_features - mean2
         cov1 = 1.0 / (n - 1) * diff1.t().mm(diff1)
-        cov2 = 1.0 / (n - 1) * diff2.t().mm(diff2)
+        cov2 = 1.0 / (m - 1) * diff2.t().mm(diff2)
 
         # compute fid
         return _compute_fid(mean1, cov1, mean2, cov2).to(orig_dtype)


Test patch:
diff --git a/tests/image/test_fid.py b/tests/image/test_fid.py
--- a/tests/image/test_fid.py
+++ b/tests/image/test_fid.py
@@ -125,21 +125,25 @@ def __len__(self):
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="test is too slow without gpu")
 @pytest.mark.skipif(not _TORCH_FIDELITY_AVAILABLE, reason="test requires torch-fidelity")
-def test_compare_fid(tmpdir, feature=2048):
+@pytest.mark.parametrize("equal_size", [False, True])
+def test_compare_fid(tmpdir, equal_size, feature=2048):
     """check that the hole pipeline give the same result as torch-fidelity."""
     from torch_fidelity import calculate_metrics
 
     metric = FrechetInceptionDistance(feature=feature).cuda()
 
+    n = 100
+    m = 100 if equal_size else 90
+
     # Generate some synthetic data
-    img1 = torch.randint(0, 180, (100, 3, 299, 299), dtype=torch.uint8)
-    img2 = torch.randint(100, 255, (100, 3, 299, 299), dtype=torch.uint8)
+    img1 = torch.randint(0, 180, (n, 3, 299, 299), dtype=torch.uint8)
+    img2 = torch.randint(100, 255, (m, 3, 299, 299), dtype=torch.uint8)
 
     batch_size = 10
-    for i in range(img1.shape[0] // batch_size):
+    for i in range(n // batch_size):
         metric.update(img1[batch_size * i : batch_size * (i + 1)].cuda(), real=True)
 
-    for i in range(img2.shape[0] // batch_size):
+    for i in range(m // batch_size):
         metric.update(img2[batch_size * i : batch_size * (i + 1)].cuda(), real=False)
 
     torch_fid = calculate_metrics(