Optionally Avoid recomputing features

torchmetrics/image/fid.py

@@ -159,6 +159,9 @@ class FID(Metric):
             - an ``nn.Module`` for using a custom feature extractor. Expects that its forward method returns
               an ``[N,d]`` matrix where ``N`` is the batch size and ``d`` is the feature size.
 
+        reset_real_features: Whether to also reset the real features. Since in many cases the real dataset does not
+            change, the features can cached them to avoid recomputing them which is costly. Set this to ``False`` if
+            your dataset does not change.
         compute_on_step:
             Forward only calls ``update()`` and return ``None`` if this is set to ``False``.
         dist_sync_on_step:
@@ -209,6 +212,7 @@ class FID(Metric):
     def __init__(
         self,
         feature: Union[int, torch.nn.Module] = 2048,
+        reset_real_features: bool = True,
         compute_on_step: bool = False,
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
@@ -248,6 +252,13 @@ def __init__(
         self.add_state("real_features", [], dist_reduce_fx=None)
         self.add_state("fake_features", [], dist_reduce_fx=None)
 
+        if reset_real_features:
+            exclude_states = ()
+        else:
+            exclude_states = ("real_features",)
+
+        self._reset_excluded_states = exclude_states
+
     def update(self, imgs: Tensor, real: bool) -> None:  # type: ignore
         """Update the state with extracted features.
 
@@ -282,3 +293,6 @@ def compute(self) -> Tensor:
 
         # compute fid
         return _compute_fid(mean1, cov1, mean2, cov2).to(orig_dtype)
+
+    def reset(self, exclude_states: Optional[Sequence[str]] = None) -> None:
+        super().reset({*self._reset_excluded_states, *exclude_states})

torchmetrics/image/kid.py

@@ -110,6 +110,9 @@ class KID(Metric):
             Scale-length of polynomial kernel. If set to ``None`` will be automatically set to the feature size
         coef:
             Bias term in the polynomial kernel.
+        reset_real_features: Whether to also reset the real features. Since in many cases the real dataset does not
+            change, the features can cached them to avoid recomputing them which is costly. Set this to ``False`` if
+            your dataset does not change.
         compute_on_step:
             Forward only calls ``update()`` and return ``None`` if this is set to ``False``.
         dist_sync_on_step:
@@ -174,6 +177,7 @@ def __init__(
         degree: int = 3,
         gamma: Optional[float] = None,  # type: ignore
         coef: float = 1.0,
+        reset_real_features: bool = True,
         compute_on_step: bool = False,
         dist_sync_on_step: bool = False,
         process_group: Optional[Any] = None,
@@ -235,6 +239,13 @@ def __init__(
         self.add_state("real_features", [], dist_reduce_fx=None)
         self.add_state("fake_features", [], dist_reduce_fx=None)
 
+        if reset_real_features:
+            exclude_states = ()
+        else:
+            exclude_states = ("real_features",)
+
+        self._reset_excluded_states = exclude_states
+
     def update(self, imgs: Tensor, real: bool) -> None:  # type: ignore
         """Update the state with extracted features.
 
@@ -276,3 +287,6 @@ def compute(self) -> Tuple[Tensor, Tensor]:
             kid_scores_.append(o)
         kid_scores = torch.stack(kid_scores_)
         return kid_scores.mean(), kid_scores.std(unbiased=False)
+
+    def reset(self, exclude_states: Optional[Sequence[str]] = None) -> None:
+        super().reset({*self._reset_excluded_states, *exclude_states})

torchmetrics/metric.py

@@ -394,7 +394,7 @@ def compute(self) -> Any:
         """Override this method to compute the final metric value from state variables synchronized across the
         distributed backend."""
 
-    def reset(self) -> None:
+    def reset(self, exclude_states: Optional[Sequence[str]] = None) -> None:
         """This method automatically resets the metric state variables to their default value."""
         self._update_called = False
         self._forward_cache = None
@@ -403,6 +403,8 @@ def reset(self) -> None:
             self._computed = None
 
         for attr, default in self._defaults.items():
+            if exclude_states is not None and attr in exclude_states:
+                continue
             current_val = getattr(self, attr)
             if isinstance(default, Tensor):
                 setattr(self, attr, default.detach().clone().to(current_val.device))