Merge pull request #227 from Adversarian/main

Robust batched FID calculation

Author: lucidrains

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -26,8 +26,7 @@
 from accelerate import Accelerator
 
 import numpy as np
-from pytorch_fid.inception import InceptionV3
-from pytorch_fid.fid_score import calculate_frechet_distance
+from denoising_diffusion_pytorch.fid_evaluation import FIDEvaluation
 
 from denoising_diffusion_pytorch.version import __version__
 
@@ -610,7 +609,7 @@ def p_mean_variance(self, x, t, x_self_cond = None, clip_denoised = True):
         model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start = x_start, x_t = x, t = t)
         return model_mean, posterior_variance, posterior_log_variance, x_start
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def p_sample(self, x, t: int, x_self_cond = None):
         b, *_, device = *x.shape, self.device
         batched_times = torch.full((b,), t, device = device, dtype = torch.long)
@@ -619,7 +618,7 @@ def p_sample(self, x, t: int, x_self_cond = None):
         pred_img = model_mean + (0.5 * model_log_variance).exp() * noise
         return pred_img, x_start
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def p_sample_loop(self, shape, return_all_timesteps = False):
         batch, device = shape[0], self.device
 
@@ -638,7 +637,7 @@ def p_sample_loop(self, shape, return_all_timesteps = False):
         ret = self.unnormalize(ret)
         return ret
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def ddim_sample(self, shape, return_all_timesteps = False):
         batch, device, total_timesteps, sampling_timesteps, eta, objective = shape[0], self.device, self.num_timesteps, self.sampling_timesteps, self.ddim_sampling_eta, self.objective
 
@@ -680,13 +679,13 @@ def ddim_sample(self, shape, return_all_timesteps = False):
         ret = self.unnormalize(ret)
         return ret
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def sample(self, batch_size = 16, return_all_timesteps = False):
         image_size, channels = self.image_size, self.channels
         sample_fn = self.p_sample_loop if not self.is_ddim_sampling else self.ddim_sample
         return sample_fn((batch_size, channels, image_size, image_size), return_all_timesteps = return_all_timesteps)
 
-    @torch.no_grad()
+    @torch.inference_mode()
     def interpolate(self, x1, x2, t = None, lam = 0.5):
         b, *_, device = *x1.shape, x1.device
         t = default(t, self.num_timesteps - 1)
@@ -738,7 +737,7 @@ def p_losses(self, x_start, t, noise = None, offset_noise_strength = None):
 
         x_self_cond = None
         if self.self_condition and random() < 0.5:
-            with torch.no_grad():
+            with torch.inference_mode():
                 x_self_cond = self.model_predictions(x, t).pred_x_start
                 x_self_cond.detach_()
 
@@ -829,7 +828,9 @@ def __init__(
         convert_image_to = None,
         calculate_fid = True,
         inception_block_idx = 2048,
-        max_grad_norm = 1.
+        max_grad_norm = 1.,
+        num_fid_samples = 50000,
+        save_best_and_latest_only = False
     ):
         super().__init__()
 
@@ -845,21 +846,15 @@ def __init__(
         self.model = diffusion_model
         self.channels = diffusion_model.channels
 
-        # InceptionV3 for fid-score computation
-
-        self.inception_v3 = None
-
-        if calculate_fid:
-            assert inception_block_idx in InceptionV3.BLOCK_INDEX_BY_DIM
-            block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[inception_block_idx]
-            self.inception_v3 = InceptionV3([block_idx])
-            self.inception_v3.to(self.device)
-
         # sampling and training hyperparameters
 
         assert has_int_squareroot(num_samples), 'number of samples must have an integer square root'
         self.num_samples = num_samples
         self.save_and_sample_every = save_and_sample_every
+        if save_best_and_latest_only:
+            assert calculate_fid, "`calculate_fid` must be True to provide a means for model evaluation for `save_best_and_latest_only`."
+            self.best_fid = 1e10 # infinite
+        self.save_best_and_latest_only = save_best_and_latest_only
 
         self.batch_size = train_batch_size
         self.gradient_accumulate_every = gradient_accumulate_every
@@ -898,6 +893,27 @@ def __init__(
 
         self.model, self.opt = self.accelerator.prepare(self.model, self.opt)
 
+        # FID-score computation
+
+        if calculate_fid:
+            self.calculate_fid = True
+            if not self.model.is_ddim_sampling:
+                self.accelerator.print(
+                    "WARNING: Robust FID computation requires a lot of generated samples and can therefore be very time consuming."\
+                    "Consider using DDIM sampling to save time."
+                )
+            self.fid_scorer = FIDEvaluation(
+                batch_size=self.batch_size,
+                dl=self.dl,
+                sampler=self.ema.ema_model,
+                channels=self.channels,
+                accelerator=self.accelerator,
+                stats_dir=results_folder,
+                device=self.device,
+                num_fid_samples=num_fid_samples,
+                inception_block_idx=inception_block_idx
+            )
+
     @property
     def device(self):
         return self.accelerator.device
@@ -937,31 +953,6 @@ def load(self, milestone):
         if exists(self.accelerator.scaler) and exists(data['scaler']):
             self.accelerator.scaler.load_state_dict(data['scaler'])
 
-    @torch.no_grad()
-    def calculate_activation_statistics(self, samples):
-        assert exists(self.inception_v3)
-
-        features = self.inception_v3(samples)[0]
-        features = rearrange(features, '... 1 1 -> ...').cpu().numpy()
-
-        mu = np.mean(features, axis = 0)
-        sigma = np.cov(features, rowvar = False)
-        return mu, sigma
-
-    def fid_score(self, real_samples, fake_samples):
-
-        if self.channels == 1:
-            real_samples, fake_samples = map(lambda t: repeat(t, 'b 1 ... -> b c ...', c = 3), (real_samples, fake_samples))
-
-        min_batch = min(real_samples.shape[0], fake_samples.shape[0])
-        real_samples, fake_samples = map(lambda t: t[:min_batch], (real_samples, fake_samples))
-
-        m1, s1 = self.calculate_activation_statistics(real_samples)
-        m2, s2 = self.calculate_activation_statistics(fake_samples)
-
-        fid_value = calculate_frechet_distance(m1, s1, m2, s2)
-        return fid_value
-
     def train(self):
         accelerator = self.accelerator
         device = accelerator.device
@@ -999,21 +990,27 @@ def train(self):
                     if self.step != 0 and self.step % self.save_and_sample_every == 0:
                         self.ema.ema_model.eval()
 
-                        with torch.no_grad():
+                        with torch.inference_mode():
                             milestone = self.step // self.save_and_sample_every
                             batches = num_to_groups(self.num_samples, self.batch_size)
                             all_images_list = list(map(lambda n: self.ema.ema_model.sample(batch_size=n), batches))
 
                         all_images = torch.cat(all_images_list, dim = 0)
 
                         utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = int(math.sqrt(self.num_samples)))
-                        self.save(milestone)
 
                         # whether to calculate fid
 
-                        if exists(self.inception_v3):
-                            fid_score = self.fid_score(real_samples = data, fake_samples = all_images)
+                        if self.calculate_fid:
+                            fid_score = self.fid_scorer.fid_score()
                             accelerator.print(f'fid_score: {fid_score}')
+                        if self.save_best_and_latest_only:
+                            if self.best_fid > fid_score:
+                                self.best_fid = fid_score
+                                self.save("best")
+                            self.save("latest")
+                        else:
+                            self.save(milestone)
 
                 pbar.update(1)
 

denoising_diffusion_pytorch/fid_evaluation.py

@@ -0,0 +1,106 @@
+import math
+import os
+
+import numpy as np
+import torch
+from einops import rearrange, repeat
+from pytorch_fid.fid_score import calculate_frechet_distance
+from pytorch_fid.inception import InceptionV3
+from torch.nn.functional import adaptive_avg_pool2d
+from tqdm.auto import tqdm
+
+
+def num_to_groups(num, divisor):
+    groups = num // divisor
+    remainder = num % divisor
+    arr = [divisor] * groups
+    if remainder > 0:
+        arr.append(remainder)
+    return arr
+
+
+class FIDEvaluation:
+    def __init__(
+        self,
+        batch_size,
+        dl,
+        sampler,
+        channels=3,
+        accelerator=None,
+        stats_dir="./results",
+        device="cuda",
+        num_fid_samples=50000,
+        inception_block_idx=2048,
+    ):
+        self.batch_size = batch_size
+        self.n_samples = num_fid_samples
+        self.device = device
+        self.channels = channels
+        self.dl = dl
+        self.sampler = sampler
+        self.stats_dir = stats_dir
+        self.print_fn = print if accelerator is None else accelerator.print
+        assert inception_block_idx in InceptionV3.BLOCK_INDEX_BY_DIM
+        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[inception_block_idx]
+        self.inception_v3 = InceptionV3([block_idx]).to(device)
+        self.dataset_stats_loaded = False
+
+    def calculate_inception_features(self, samples):
+        if self.channels == 1:
+            samples = repeat(samples, "b 1 ... -> b c ...", c=3)
+        features = self.inception_v3(samples)[0]
+        if features.size(2) != 1 or features.size(3) != 1:
+            features = adaptive_avg_pool2d(features, output_size=(1, 1))
+        features = rearrange(features, "... 1 1 -> ...")
+        return features
+
+    def load_or_precalc_dataset_stats(self):
+        path = os.path.join(self.stats_dir, "dataset_stats")
+        try:
+            ckpt = np.load(path + ".npz")
+            self.m2, self.s2 = ckpt["m2"], ckpt["s2"]
+            self.print_fn("Dataset stats loaded from disk.")
+            ckpt.close()
+        except OSError:
+            num_batches = int(math.ceil(self.n_samples / self.batch_size))
+            stacked_real_features = []
+            self.print_fn(
+                f"Stacking Inception features for {self.n_samples} samples from the real dataset."
+            )
+            for _ in tqdm(range(num_batches)):
+                try:
+                    real_samples = next(self.dl)
+                except StopIteration:
+                    break
+                real_samples = real_samples.to(self.device)
+                real_features = self.calculate_inception_features(real_samples)
+                stacked_real_features.append(real_features)
+            stacked_real_features = (
+                torch.cat(stacked_real_features, dim=0).cpu().numpy()
+            )
+            m2 = np.mean(stacked_real_features, axis=0)
+            s2 = np.cov(stacked_real_features, rowvar=False)
+            np.savez_compressed(path, m2=m2, s2=s2)
+            self.print_fn(f"Dataset stats cached to {path}.npz for future use.")
+            self.m2, self.s2 = m2, s2
+        self.dataset_stats_loaded = True
+
+    @torch.inference_mode()
+    def fid_score(self):
+        if not self.dataset_stats_loaded:
+            self.load_or_precalc_dataset_stats()
+        self.sampler.eval()
+        batches = num_to_groups(self.n_samples, self.batch_size)
+        stacked_fake_features = []
+        self.print_fn(
+            f"Stacking Inception features for {self.n_samples} generated samples."
+        )
+        for batch in tqdm(batches):
+            fake_samples = self.sampler.sample(batch_size=batch)
+            fake_features = self.calculate_inception_features(fake_samples)
+            stacked_fake_features.append(fake_features)
+        stacked_fake_features = torch.cat(stacked_fake_features, dim=0).cpu().numpy()
+        m1 = np.mean(stacked_fake_features, axis=0)
+        s1 = np.cov(stacked_fake_features, rowvar=False)
+
+        return calculate_frechet_distance(m1, s1, self.m2, self.s2)