fid calculation

.gitignore

@@ -1,3 +1,13 @@
 .DS_Store
 __pycache__/
 
+ckpt/
+*.ipynb
+datasets/cifar_test
+datasets/cifar_train
+.venv/
+anaconda.sh
+improved_diffusion.egg-info/
+wandb/
+*.pt
+results/

improved_diffusion/fid_evaluation.py

@@ -0,0 +1,137 @@
+"""
+Code from https://github.com/lucidrains/denoising-diffusion-pytorch/blob/main/denoising_diffusion_pytorch/fid_evaluation.py
+"""
+
+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
+
+from datetime import datetime
+
+
+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=10000,
+        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
+        self.samples_path = None  # 샘플 경로를 저장할 새 속성
+
+        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 set_samples_path(self, path):
+        """생성된 샘플의 경로를 설정하는 메서드"""
+        self.samples_path = path
+
+    def calculate_inception_features(self, samples):
+        if self.channels == 1:
+            samples = repeat(samples, "b 1 ... -> b c ...", c=3)
+
+        self.inception_v3.eval()
+        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, f"{datetime.now().strftime('%Y%m%d_%H%M')}_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
+                if isinstance(real_samples, (list, tuple)):
+                    real_samples = real_samples[0]
+                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)
+        if self.samples_path is None:
+            raise ValueError("Samples path is not set. Call set_samples_path() before calculating FID score.")
+
+        # 저장된 샘플 로드
+        samples_data = np.load(self.samples_path)
+        if samples_data.files[0] == "arr_0":
+            samples = samples_data["arr_0"]
+        else:
+            samples = samples_data[samples_data.files[0]]
+
+        if len(samples) < self.n_samples:
+            raise ValueError(f"Not enough samples. Found {len(samples)}, need {self.n_samples}")
+
+        samples = samples[: self.n_samples]  # 필요한 만큼만 사용
+
+        stacked_fake_features = []
+        self.print_fn(f"Stacking Inception features for {self.n_samples} generated samples.")
+        # for batch in tqdm(batches):
+        for i in tqdm(range(0, self.n_samples, self.batch_size)):
+            batch = samples[i : i + self.batch_size]
+            batch = torch.from_numpy(batch).float().to(self.device) / 127.5 - 1  # 정규화
+            batch = batch.permute(0, 3, 1, 2)  # NHWC to NCHW
+            fake_features = self.calculate_inception_features(batch)
+            # 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)