kl fid and attention

2 files changed, 127 insertions, 82 deletions

diff --git a/main.py b/main.py
index aa07c3d..81f832d 100644
--- a/main.py
+++ b/main.py
@@ -3,6 +3,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torch.nn.attention as attention
 from torch.utils.data import DataLoader, Dataset
+from rich.progress import track
 
 from torchvision import datasets, transforms
 
@@ -14,7 +15,7 @@ import cv2
 
 from trainer import Trainer
 from autoencoder import Autoencoder
-from utils import fid
+from utils import *
 
 
 class SelfAttention(nn.Module):
@@ -303,104 +304,104 @@ if __name__ == '__main__':
     ##############
 
     with torch.no_grad():
-        # # Forward diffusion
-        # img, label = dataset[np.random.randint(0, len(dataset))]
-        # img = img.to(DEVICE)
-        # if autoencoder is not None:
-        #     img = autoencoder.encode(img.unsqueeze(0)).squeeze(0)
-        # img = img * 2 - 1
+        # Forward diffusion
+        img, label = dataset[np.random.randint(0, len(dataset))]
+        img = img.to(DEVICE)
+        if autoencoder is not None:
+            img = autoencoder.encode(img.unsqueeze(0)).squeeze(0)
+        img = img * 2 - 1
 
-        # nb_plots = 10
-        # plots_id = [i for i in np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)]
+        nb_plots = 10
+        plots_id = [i for i in np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)]
 
-        # xs = forward_diffusion(img)
+        xs = forward_diffusion(img)
 
-        # plt.figure(figsize=(nb_plots, 2.5))
-        # for plot_i, t in enumerate(plots_id):
-        #     x = xs[t]
-        #     plt.subplot(2, nb_plots + 1, plot_i + 2)
-        #     plt.title(f"t={t}")
-        #     plt.imshow(tensor_to_image(x), interpolation='none')
-        #     plt.axis("off")
+        plt.figure(figsize=(nb_plots, 2.5))
+        for plot_i, t in enumerate(plots_id):
+            x = xs[t]
+            plt.subplot(2, nb_plots + 1, plot_i + 2)
+            plt.title(f"t={t}")
+            plt.imshow(tensor_to_image(x), interpolation='none')
+            plt.axis("off")
 
-        # for plot_i, t in enumerate(plots_id):
-        #     x = q_xt_x0(img, t)[0]
-        #     plt.subplot(2, nb_plots + 1, nb_plots + plot_i + 3)
-        #     plt.imshow(tensor_to_image(x), interpolation='none')
-        #     plt.axis("off")
+        for plot_i, t in enumerate(plots_id):
+            x = q_xt_x0(img, t)[0]
+            plt.subplot(2, nb_plots + 1, nb_plots + plot_i + 3)
+            plt.imshow(tensor_to_image(x), interpolation='none')
+            plt.axis("off")
 
-        # plt.subplot(2, nb_plots + 1, 1)
-        # plt.text(0, 0.5, "Implicit", fontsize=12)
-        # plt.axis("off")
-        # plt.subplot(2, nb_plots + 1, nb_plots + 2)
-        # plt.text(0, 0.5, "Explicit", fontsize=12)
-        # plt.axis("off")
+        plt.subplot(2, nb_plots + 1, 1)
+        plt.text(0, 0.5, "Implicit", fontsize=12)
+        plt.axis("off")
+        plt.subplot(2, nb_plots + 1, nb_plots + 2)
+        plt.text(0, 0.5, "Explicit", fontsize=12)
+        plt.axis("off")
 
-        # plt.suptitle("Forward diffusion")
-        # plt.tight_layout()
-        # plt.savefig("forward_diffusion.tmp.png")
+        plt.suptitle("Forward diffusion")
+        plt.tight_layout()
+        plt.savefig("forward_diffusion.tmp.png")
 
 
-        # # Backward diffusion
-        # t_plots = np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)
+        # Backward diffusion
+        t_plots = np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)
 
-        # n_classes = 10
+        n_classes = 10
 
-        # x = torch.randn(n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE)
-        # vec = torch.tensor([[min(i, NB_LABEL-1)] for i in range(n_classes)], dtype=torch.int64)
-        # vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32)
+        x = torch.randn(n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE, device=DEVICE)
+        vec = torch.tensor([[min(i, NB_LABEL-1)] for i in range(n_classes)], dtype=torch.int64)
+        vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32)
 
-        # plt.figure(figsize=(nb_plots, n_classes))
-        # plt.suptitle("Backward diffusion")
-        # for t in range(DIFFU_STEPS, 0, -1):
-        #     t_tensor = torch.tensor([[t]] * n_classes, device=DEVICE, dtype=torch.float32)
-        #     x = p_xt_1_xt(model, x, t_tensor, vec)
-        #     if t in t_plots:
-        #         t_plot_i = nb_plots - t_plots.tolist().index(t) - 1
-        #         for class_i in range(n_classes):
-        #             plt.subplot(n_classes, nb_plots, t_plot_i + nb_plots * class_i + 1)
-        #             if class_i == 0:
-        #                 plt.title(f"t={t}")
-        #             plt.imshow(tensor_to_image(x[class_i]))
-        #             plt.axis("off")
-        # plt.tight_layout()
-        # plt.savefig("backward_diffusion.tmp.png")
+        plt.figure(figsize=(nb_plots, n_classes))
+        plt.suptitle("Backward diffusion")
+        for t in track(range(DIFFU_STEPS, 0, -1), description='Diffusing...'):
+            t_tensor = torch.tensor([[t]] * n_classes, device=DEVICE, dtype=torch.float32)
+            x = p_xt_1_xt(model, x, t_tensor, vec)
+            if t in t_plots:
+                t_plot_i = nb_plots - t_plots.tolist().index(t) - 1
+                for class_i in range(n_classes):
+                    plt.subplot(n_classes, nb_plots, t_plot_i + nb_plots * class_i + 1)
+                    if class_i == 0:
+                        plt.title(f"t={t}")
+                    plt.imshow(tensor_to_image(x[class_i]))
+                    plt.axis("off")
+        plt.tight_layout()
+        plt.savefig("backward_diffusion.tmp.png")
 
 
-        # # Benchmark
-        # x = torch.randn(nb_plots * n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE).to(DEVICE)
-        # vec = sum([[[min(i, NB_LABEL-1)]] * nb_plots for i in range(n_classes)], [])
-        # vec = torch.tensor(vec, device=DEVICE, dtype=torch.int64)
-        # vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32)
+        # Benchmark
+        x = torch.randn(nb_plots * n_classes, NB_CHANNEL, IMG_SIZE, IMG_SIZE).to(DEVICE)
+        vec = sum([[[min(i, NB_LABEL-1)]] * nb_plots for i in range(n_classes)], [])
+        vec = torch.tensor(vec, device=DEVICE, dtype=torch.int64)
+        vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32)
 
-        # for ti in range(DIFFU_STEPS, 0, -1):
-        #     t = torch.tensor([[ti]] * n_classes * nb_plots, device=DEVICE, dtype=torch.float32)
-        #     x = p_xt_1_xt(model, x, t, vec)
+        for ti in track(range(DIFFU_STEPS, 0, -1), description='Benchmarking...'):
+            t = torch.tensor([[ti]] * n_classes * nb_plots, device=DEVICE, dtype=torch.float32)
+            x = p_xt_1_xt(model, x, t, vec)
 
-        # x = x * 0.5 + 0.5
-        # x = x.clamp(0, 1)
+        x = x * 0.5 + 0.5
+        x = x.clamp(0, 1)
 
-        # plt.figure(figsize=(nb_plots, n_classes))
-        # for i in range(nb_plots):
-        #     for j in range(n_classes):
-        #         id = i * n_classes + j
-        #         img = x[id]
-        #         if autoencoder is not None:
-        #             img = train_dataset.autoencoder.decode(img.unsqueeze(0)).squeeze(0)
-        #         plt.subplot(n_classes, nb_plots, id + 1)
-        #         plt.imshow(tensor_to_image(img))
-        #         plt.axis("off")
-        # plt.tight_layout()
-        # plt.savefig("benchmark.tmp.png")
+        plt.figure(figsize=(nb_plots, n_classes))
+        for i in range(nb_plots):
+            for j in range(n_classes):
+                id = i * n_classes + j
+                img = x[id]
+                if autoencoder is not None:
+                    img = train_dataset.autoencoder.decode(img.unsqueeze(0)).squeeze(0)
+                plt.subplot(n_classes, nb_plots, id + 1)
+                plt.imshow(tensor_to_image(img))
+                plt.axis("off")
+        plt.tight_layout()
+        plt.savefig("benchmark.tmp.png")
 
 
-        # FID
-        n_samples = 10
-        batch_size = 4
-        n_batches = n_samples // batch_size
+        # Metrics
+        batch_size = 64
+        n_batches = 16
+        n_samples = batch_size * n_batches
 
         fakes = np.zeros((0, NB_CHANNEL, IMG_SIZE, IMG_SIZE))
-        for _ in range(n_batches):
+        for _ in track(range(n_batches), description=f'Sampling {n_samples} images...'):
             x = torch.randn(batch_size, NB_CHANNEL, IMG_SIZE, IMG_SIZE).to(DEVICE)
             vec = torch.randint(0, NB_LABEL, (batch_size,)).to(DEVICE)
             vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32)
@@ -414,8 +415,13 @@ if __name__ == '__main__':
         reals = torch.stack([dataset[i][0] for i in range(n_samples)]).cpu().numpy()
         reals = reals * 2 - 1
 
+
         fid_score = fid(reals, fakes)
         print(f"FID score: {fid_score}")
 
 
+        kl_score = kl(reals, fakes)
+        print(f"KL divergence: {kl_score}")
+
+
         # plt.show()
diff --git a/utils.py b/utils.py
index 19a0426..4e680a8 100644
--- a/utils.py
+++ b/utils.py
@@ -1,4 +1,6 @@
 import numpy as np
+import scipy.linalg
+import matplotlib.pyplot as plt
 
 
 def fid(reals, fakes):
@@ -8,7 +10,6 @@ def fid(reals, fakes):
         reals (numpy.array): Real images.
         fakes (numpy.array): Fake images.
     """
-    print(reals.shape, fakes.shape)
     reals = reals.reshape(reals.shape[0], -1)
     fakes = fakes.reshape(fakes.shape[0], -1)
 
@@ -19,7 +20,45 @@ def fid(reals, fakes):
 
     diff = mu_real - mu_fake
     covmean = np.dot(sigma_real, sigma_fake.T)
-    covmean = np.sqrt(covmean * (covmean > 0))
-    print(np.trace(covmean))
+    covmean, _ = scipy.linalg.sqrtm(sigma_real.dot(sigma_fake), disp=False)
+
+    if not np.isfinite(covmean).all():
+        eps=1e-6
+        offset = np.eye(sigma_real.shape[0]) * eps
+        ncovmean = scipy.linalg.sqrtm((sigma_real + offset).dot(sigma_fake + offset))
+        covmean = ncovmean
+
+    if np.iscomplexobj(covmean):
+        covmean = covmean.real
 
     return diff @ diff + np.trace(sigma_real) + np.trace(sigma_fake) - 2 * np.trace(covmean)
+
+
+def kl(reals, fakes):
+    """KL divergence calculation.
+
+    Args:
+        reals (numpy.array): Real images.
+        fakes (numpy.array): Fake images.
+    """
+    reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1)
+    fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1)
+
+    hist_real = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals)
+    hist_fake = np.apply_along_axis(lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes)
+
+    plt.figure()
+    colors = ['#ff0000', '#00ff00', '#0000ff']
+    for i in range(reals.shape[0]):
+        plt.plot(hist_real[i], label='real', color=colors[i])
+        plt.plot(hist_fake[i], label='fake', color=colors[i], linestyle='dashed')
+    plt.legend()
+    plt.savefig('hist.tmp.png')
+
+    hist_real = hist_real + 1
+    hist_fake = hist_fake + 1
+
+    hist_real = hist_real / np.sum(hist_real)
+    hist_fake = hist_fake / np.sum(hist_fake)
+
+    return np.mean(np.log(hist_real / hist_fake))
\ No newline at end of file