From 47b3bc6fbfed3360e7526dcf5275f9175e0fde10 Mon Sep 17 00:00:00 2001 From: gdamms Date: Thu, 20 Jun 2024 15:59:26 +0200 Subject: kl fid and attention --- main.py | 200 ++++++++++++++++++++++++++++++++------------------------------- utils.py | 45 +++++++++++++- 2 files changed, 145 insertions(+), 100 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 - - # nb_plots = 10 - # plots_id = [i for i in np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)] - - # 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") - - # 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.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) - - # 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) - - # 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") - - - # # 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) - - # 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") - - - # FID - n_samples = 10 - batch_size = 4 - n_batches = n_samples // batch_size + # 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)] + + 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") + + 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.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) + + 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) + + 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) + + 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) + + 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") + + + # 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 -- cgit v1.3.1