How to calculate FID score?

Use multi-GPU to calc FID score · Lightning-AI/pytorch-lightning · Discussion #17302 · GitHub

origin question:
I’m training the diffusion model and I want to calculate the FID score for each some steps.
Because the FID calculation needs to generate so many samples, e.g. 5000, I hope to distribute them to each GPU.
I know I can create a Callback class and implement it in on_validation_epoch_start : generate 5000 samples, get the result, and log it, but it’s on a single GPU.

my question:
I want to eval the FID score using GPU or GPUs on val_epoch(generate 5000 samples, get the result), can you offer me some examples, thanks.

@awaelchli @SkafteNicki @jirka

from pytorch_lightning.utilities import CombinedLoader

class ImgGenerator(Dataset):
    def __init__(self, generate_num=50000, generate_size=(3, 32, 32)):
        self.generate_num = generate_num
        self.generate_size = generate_size

    def __len__(self):
        return self.generate_num

    def __getitem__(self, idx):
        generate_item = torch.zeros(self.generate_size)
        return generate_item

# 50K validation
val_iterables = {
    'pred_images': DataLoader(ImgGenerator(generate_num=50000, generate_size=(3, 32, 32)),
                              shuffle=False, drop_last=False,
                              batch_size=self.batch_size, pin_memory=self.pin_mem,
                              num_workers=self.num_workers,
                              persistent_workers=True),
    'true_images': DataLoader(self.test_dataset, shuffle=False, drop_last=False,
                              batch_size=self.batch_size, pin_memory=self.pin_mem,
                              num_workers=self.num_workers,
                              persistent_workers=True),
}
val_combined_loader = CombinedLoader(val_iterables, 'max_size')