Allow caching on cpu when compute some metrics.

[kynk94]

🚀 Feature

Currently, some metrics like FID and KID only use GPU computing.
Since all datasets are cached in gpu memory, there is a big limitation on hardware.
It may be helpful to add compute_on_cpu option to compute these metrics like below.

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)

    if self.compute_on_cpu:
        features = features.detach().cpu()

    if real:
        self.real_features.append(features)
    else:
        self.fake_features.append(features)

Motivation

I can compute FID on cloud system but not my local, because of the memory issue.
Currently, warn large memory foorprint. #468

[github-actions]

Hi! thanks for your contribution!, great first issue!

[SkafteNicki]

Hi @kynk94,
I do not completely understand the problem here. If you initialize a metric, it will by default be on the CPU:

Do you want to keep the overall metric on GPU, meaning that features = self.inception(imgs) is still calculated on the GPU but when you store the features for later you want that to be on CPU?

[Borda]

My understanding is that if you run training and you use all GPU memory it could be interesting to force/offload the metric computation to CPU... is it correct @kynk94?

[kynk94]

Yes, that's right.
If training on GPU, the inception features are stored in GPU memory to compute metric.
If have a lot of GPU, there are enough GPU memory.
But most local desktops not have enough GPU memory, and CPU memory capacity is larger than GPU memory.

In this case, to avoid memory overflow errors, need to cache the features on CPU memory and compute metric on CPU