Repository: Lightning-AI/torchmetrics

Issue:
Change FID to be online
## 🚀 Feature

The idea is to compute mean and covariance online after each update instead of storing the features and then computing the mean and covariance

### Motivation

The motivation is to avoid storing the features which will lead to a way more memory efficient metric.

### Pitch

This can be done the following way:
```python
def update(self, imgs: Tensor, real: bool) -> None:  # type: ignore
        """Update the state with extracted features.

        Args:
            imgs: tensor with images feed to the feature extractor
            real: bool indicating if ``imgs`` belong to the real or the fake distribution
        """
        features = self.inception(imgs).double()
        self.orig_dtype = features.dtype

        if real:
            self.real_features_sum += features.sum(dim=0)
            self.real_features_cov_sum += features.t().mm(features)
            self.real_features_num_samples += imgs.shape[0]
        else:
            self.fake_features_sum += features.sum(dim=0)
            self.fake_features_cov_sum += features.t().mm(features)
            self.fake_features_num_samples += imgs.shape[0]

    def compute(self) -> Tensor:
        """Calculate FID score based on accumulated extracted features from the two distributions."""
        mean_real = (self.real_features_sum / self.real_features_num_samples).unsqueeze(0)
        mean_fake = (self.fake_features_sum / self.fake_features_num_samples).unsqueeze(dim=0)

        cov_real = (self.real_features_cov_sum / self.real_features_num_samples - mean_real.t().mm(mean_real)) / (
            self.real_features_num_samples - 1
        )
        cov_fake = (self.fake_features_cov_sum - self.fake_features_num_samples * mean_fake.t().mm(mean_fake)) / (
            self.fake_features_num_samples - 1
        )
        return _compute_fid(mean_real, cov_real, mean_fake, cov_fake).to(self.orig_dtype)
```



Discussion:

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
@@ -11,6 +11,7 @@
 # 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 copy import deepcopy
 from typing import Any, List, Optional, Union
 
 import numpy as np
@@ -20,8 +21,7 @@
 from torch.nn import Module
 
 from torchmetrics.metric import Metric
-from torchmetrics.utilities import rank_zero_info, rank_zero_warn
-from torchmetrics.utilities.data import dim_zero_cat
+from torchmetrics.utilities import rank_zero_info
 from torchmetrics.utilities.imports import _SCIPY_AVAILABLE, _TORCH_FIDELITY_AVAILABLE
 
 if _TORCH_FIDELITY_AVAILABLE:
@@ -201,8 +201,13 @@ class FrechetInceptionDistance(Metric):
     is_differentiable: bool = False
     full_state_update: bool = False
 
-    real_features: List[Tensor]
-    fake_features: List[Tensor]
+    real_features_sum: Tensor
+    real_features_cov_sum: Tensor
+    real_features_num_samples: Tensor
+
+    fake_features_sum: Tensor
+    fake_features_cov_sum: Tensor
+    fake_features_num_samples: Tensor
 
     def __init__(
         self,
@@ -212,13 +217,8 @@ def __init__(
     ) -> None:
         super().__init__(**kwargs)
 
-        rank_zero_warn(
-            "Metric `FrechetInceptionDistance` will save all extracted features in buffer."
-            " For large datasets this may lead to large memory footprint.",
-            UserWarning,
-        )
-
         if isinstance(feature, int):
+            num_features = feature
             if not _TORCH_FIDELITY_AVAILABLE:
                 raise ModuleNotFoundError(
                     "FrechetInceptionDistance metric requires that `Torch-fidelity` is installed."
@@ -231,8 +231,11 @@ def __init__(
                 )
 
             self.inception = NoTrainInceptionV3(name="inception-v3-compat", features_list=[str(feature)])
+
         elif isinstance(feature, Module):
             self.inception = feature
+            dummy_image = torch.randint(0, 255, (1, 3, 299, 299), dtype=torch.uint8)
+            num_features = self.inception(dummy_image).shape[-1]
         else:
             raise TypeError("Got unknown input to argument `feature`")
 
@@ -240,8 +243,14 @@ def __init__(
             raise ValueError("Argument `reset_real_features` expected to be a bool")
         self.reset_real_features = reset_real_features
 
-        self.add_state("real_features", [], dist_reduce_fx=None)
-        self.add_state("fake_features", [], dist_reduce_fx=None)
+        mx_nb_feets = (num_features, num_features)
+        self.add_state("real_features_sum", torch.zeros(num_features).double(), dist_reduce_fx="sum")
+        self.add_state("real_features_cov_sum", torch.zeros(mx_nb_feets).double(), dist_reduce_fx="sum")
+        self.add_state("real_features_num_samples", torch.tensor(0).long(), dist_reduce_fx="sum")
+
+        self.add_state("fake_features_sum", torch.zeros(num_features).double(), dist_reduce_fx="sum")
+        self.add_state("fake_features_cov_sum", torch.zeros(mx_nb_feets).double(), dist_reduce_fx="sum")
+        self.add_state("fake_features_num_samples", torch.tensor(0).long(), dist_reduce_fx="sum")
 
     def update(self, imgs: Tensor, real: bool) -> None:  # type: ignore
         """Update the state with extracted features.
@@ -251,39 +260,39 @@ def update(self, imgs: Tensor, real: bool) -> None:  # type: ignore
             real: bool indicating if ``imgs`` belong to the real or the fake distribution
         """
         features = self.inception(imgs)
+        self.orig_dtype = features.dtype
+        features = features.double()
 
+        if features.dim() == 1:
+            features = features.unsqueeze(0)
         if real:
-            self.real_features.append(features)
+            self.real_features_sum += features.sum(dim=0)
+            self.real_features_cov_sum += features.t().mm(features)
+            self.real_features_num_samples += imgs.shape[0]
         else:
-            self.fake_features.append(features)
+            self.fake_features_sum += features.sum(dim=0)
+            self.fake_features_cov_sum += features.t().mm(features)
+            self.fake_features_num_samples += imgs.shape[0]
 
     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]
-        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)
-
-        # compute fid
-        return _compute_fid(mean1, cov1, mean2, cov2).to(orig_dtype)
+        mean_real = (self.real_features_sum / self.real_features_num_samples).unsqueeze(0)
+        mean_fake = (self.fake_features_sum / self.fake_features_num_samples).unsqueeze(0)
+
+        cov_real_num = self.real_features_cov_sum - self.real_features_num_samples * mean_real.t().mm(mean_real)
+        cov_real = cov_real_num / (self.real_features_num_samples - 1)
+        cov_fake_num = self.fake_features_cov_sum - self.fake_features_num_samples * mean_fake.t().mm(mean_fake)
+        cov_fake = cov_fake_num / (self.fake_features_num_samples - 1)
+        return _compute_fid(mean_real.squeeze(0), cov_real, mean_fake.squeeze(0), cov_fake).to(self.orig_dtype)
 
     def reset(self) -> None:
         if not self.reset_real_features:
-            # remove temporarily to avoid resetting
-            value = self._defaults.pop("real_features")
+            real_features_sum = deepcopy(self.real_features_sum)
+            real_features_cov_sum = deepcopy(self.real_features_cov_sum)
+            real_features_num_samples = deepcopy(self.real_features_num_samples)
             super().reset()
-            self._defaults["real_features"] = value
+            self.real_features_sum = real_features_sum
+            self.real_features_cov_sum = real_features_cov_sum
+            self.real_features_num_samples = real_features_num_samples
         else:
             super().reset()


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
@@ -72,13 +72,6 @@ def test_fid_pickle():
 
 def test_fid_raises_errors_and_warnings():
     """Test that expected warnings and errors are raised."""
-    with pytest.warns(
-        UserWarning,
-        match="Metric `FrechetInceptionDistance` will save all extracted features in buffer."
-        " For large datasets this may lead to large memory footprint.",
-    ):
-        _ = FrechetInceptionDistance()
-
     if _TORCH_FIDELITY_AVAILABLE:
         with pytest.raises(ValueError, match="Integer input to argument `feature` must be one of .*"):
             _ = FrechetInceptionDistance(feature=2)
@@ -106,7 +99,9 @@ def test_fid_same_input(feature):
         metric.update(img, real=True)
         metric.update(img, real=False)
 
-    assert torch.allclose(torch.cat(metric.real_features, dim=0), torch.cat(metric.fake_features, dim=0))
+    assert torch.allclose(metric.real_features_sum, metric.fake_features_sum)
+    assert torch.allclose(metric.real_features_cov_sum, metric.fake_features_cov_sum)
+    assert torch.allclose(metric.real_features_num_samples, metric.fake_features_num_samples)
 
     val = metric.compute()
     assert torch.allclose(val, torch.zeros_like(val), atol=1e-3)
@@ -167,19 +162,22 @@ def test_reset_real_features_arg(reset_real_features):
     metric.update(torch.randint(0, 180, (2, 3, 299, 299), dtype=torch.uint8), real=True)
     metric.update(torch.randint(0, 180, (2, 3, 299, 299), dtype=torch.uint8), real=False)
 
-    assert len(metric.real_features) == 1
-    assert list(metric.real_features[0].shape) == [2, 64]
+    assert metric.real_features_num_samples == 2
+    assert metric.real_features_sum.shape == torch.Size([64])
+    assert metric.real_features_cov_sum.shape == torch.Size([64, 64])
 
-    assert len(metric.fake_features) == 1
-    assert list(metric.fake_features[0].shape) == [2, 64]
+    assert metric.fake_features_num_samples == 2
+    assert metric.fake_features_sum.shape == torch.Size([64])
+    assert metric.fake_features_cov_sum.shape == torch.Size([64, 64])
 
     metric.reset()
 
     # fake features should always reset
-    assert len(metric.fake_features) == 0
+    assert metric.fake_features_num_samples == 0
 
     if reset_real_features:
-        assert len(metric.real_features) == 0
+        assert metric.real_features_num_samples == 0
     else:
-        assert len(metric.real_features) == 1
-        assert list(metric.real_features[0].shape) == [2, 64]
+        assert metric.real_features_num_samples == 2
+        assert metric.real_features_sum.shape == torch.Size([64])
+        assert metric.real_features_cov_sum.shape == torch.Size([64, 64])