From 21759fef8826367f8e08bfc1c8118957730cdfdd Mon Sep 17 00:00:00 2001 From: gdamms Date: Fri, 17 May 2024 15:38:32 +0200 Subject: more plots --- main.py | 22 +++++++++-------- plots.py | 84 +++++++++++++++++++++++++++++++++++++++++++++++----------------- 2 files changed, 74 insertions(+), 32 deletions(-) diff --git a/main.py b/main.py index 5356aed..5102fc7 100644 --- a/main.py +++ b/main.py @@ -191,7 +191,6 @@ def forward_diffusion(x0): def tensor_to_image(tensor): img = tensor.clone().detach().cpu().numpy().transpose(1, 2, 0) - img = img / 2 + 0.5 img -= img.min() img /= img.max() return img @@ -205,12 +204,12 @@ BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA)) ALPHA = 1 - BETA ALPHA_BAR = torch.cumprod(ALPHA, dim=0) -# dataset = datasets.MNIST( -# root="./data", -# train=True, -# download=True, -# transform=transforms.ToTensor(), -# ) +dataset = datasets.MNIST( + root="./data", + train=True, + download=True, + transform=transforms.ToTensor(), +) # dataset = datasets.LFWPeople( # root="./data", # download=True, @@ -219,13 +218,13 @@ ALPHA_BAR = torch.cumprod(ALPHA, dim=0) # transforms.ToTensor(), # ]), # ) -dataset = LFWcrop() +# dataset = LFWcrop() img = dataset[0][0] NB_CHANNEL, IMG_SIZE, _ = img.shape -NB_LABEL = 1 +NB_LABEL = 10 -EPOCHS = 100 +EPOCHS = 10 LEARNING_RATE = 2e-4 @@ -344,6 +343,9 @@ if __name__ == '__main__': 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): diff --git a/plots.py b/plots.py index dd4706b..db5f009 100644 --- a/plots.py +++ b/plots.py @@ -2,28 +2,19 @@ import matplotlib.pyplot as plt from matplotlib.gridspec import GridSpec import numpy as np from torchvision import datasets +import torch +import os +from rich.progress import track +from main import UNet, q_xt_xt_1, p_xt_1_xt -def q_xt_x0(x0, t): - alpha_bar = ALPHA_BAR[t] - mean = np.sqrt(alpha_bar) * x0 - std = np.sqrt(1 - alpha_bar) - eps = np.random.normal(0, 1, x0.shape) - xt = mean + std * eps +os.makedirs('plots', exist_ok=True) +os.makedirs('plots/diffusion', exist_ok=True) +os.makedirs('plots/diffusion_inverse', exist_ok=True) - return xt, eps - -def q_xt_xt_1(xt_1, t): - alpha = ALPHA[t] - mean = np.sqrt(alpha) * xt_1 - std = np.sqrt(1 - alpha) - - eps = np.random.normal(0, 1, xt_1.shape) - xt = mean + std * eps - - return xt, eps +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') DIFFU_STEPS = 1000 BETA = np.linspace(1e-4, 2e-2, DIFFU_STEPS) @@ -46,6 +37,7 @@ plt.savefig('plots/alpha_beta.tmp.png') mnist = datasets.MNIST('data', train=True, download=True) img = mnist.data[np.random.randint(0, len(mnist))].numpy() / 255 +img = img * 2 - 1 plt.figure() plt.imshow(img, cmap='gray') @@ -68,24 +60,72 @@ x_norm = np.linspace(BIN_MIN, BIN_MAX, 100) y_norm = norm_dist(x_norm, 0, 1) * 28**2 / NB_BINS * (BIN_MAX - BIN_MIN) fig = plt.figure(figsize=(10, 5)) +fig.suptitle('Diffusion naturelle') gs = GridSpec(1, 3, figure=fig) ax1 = fig.add_subplot(gs[0, 0]) ax2 = fig.add_subplot(gs[0, 1:]) -xt = img -for t in range(DIFFU_STEPS): +plots_to_save = np.linspace(1, DIFFU_STEPS, 100).astype(int) + +xt = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0) +for t in track(range(1, DIFFU_STEPS+1)): xt, eps = q_xt_xt_1(xt, t) + if t not in plots_to_save: + continue + + xt_numpy = xt.cpu().detach().numpy()[0, 0] + ax1.clear() - ax1.imshow(xt, cmap='gray') - ax1.set_title(f'xt at t={t}') + ax1.imshow(xt_numpy, cmap='gray') + ax1.set_title(f'xt at t={t:04d}') ax1.axis('off') ax2.clear() - ax2.hist(xt.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) + ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') + ax2.legend() ax2.set_yscale('log') ax2.set_ylim(y_norm.min(), 1e3) ax2.set_title(f'xt histogram') fig.savefig(f'plots/diffusion/{t:04d}.tmp.png') +os.system('convert --delay 20 --repeat 0 plots/diffusion/*.png plots/diffusion.tmp.gif') + + + +model = UNet().to(DEVICE) +model.load_state_dict(torch.load('model.pth')) + +fig = plt.figure(figsize=(10, 5)) +fig.suptitle('Diffusion inverse') +gs = GridSpec(1, 3, figure=fig) +ax1 = fig.add_subplot(gs[0, 0]) +ax2 = fig.add_subplot(gs[0, 1:]) + +xt = torch.randn(1, 1, 28, 28, device=DEVICE) +vec = torch.zeros(1, 10).to(DEVICE) +for t in track(range(DIFFU_STEPS, 0, -1)): + t_tensor = torch.tensor([[t]], device=DEVICE, dtype=torch.float32) + xt = p_xt_1_xt(model, xt, t_tensor, vec) + + if t not in plots_to_save: + continue + + xt_numpy = xt.cpu().detach().numpy()[0, 0] + + ax1.clear() + ax1.imshow(xt_numpy, cmap='gray') + ax1.set_title(f'xt at t={t:04d}') + ax1.axis('off') + + ax2.clear() + ax2.hist(xt_numpy.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX)) + ax2.plot(x_norm, y_norm, color='red', label='N(0, 1)') + ax2.legend() + ax2.set_yscale('log') + ax2.set_ylim(y_norm.min(), 1e3) + ax2.set_title(f'xt histogram') + + fig.savefig(f'plots/diffusion_inverse/{t:04d}.tmp.png') +os.system('convert --delay 20 --repeat 0 -reverse plots/diffusion_inverse/*.png plots/diffusion_inverse.tmp.gif') -- cgit v1.3.1