aboutsummaryrefslogtreecommitdiff
path: root/plots.py
diff options
context:
space:
mode:
Diffstat (limited to 'plots.py')
-rw-r--r--plots.py174
1 files changed, 0 insertions, 174 deletions
diff --git a/plots.py b/plots.py
deleted file mode 100644
index 0513495..0000000
--- a/plots.py
+++ /dev/null
@@ -1,174 +0,0 @@
-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, tensor_to_image
-from autoencoder import Autoencoder
-
-
-os.makedirs('plots', exist_ok=True)
-os.makedirs('plots/diffusion', exist_ok=True)
-os.makedirs('plots/diffusion_inverse', exist_ok=True)
-
-
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-DIFFU_STEPS = 1000
-BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE)
-BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA))
-ALPHA = 1 - BETA
-ALPHA_BAR = torch.cumprod(ALPHA, dim=0)
-
-NB_BINS = 50
-BIN_MIN = -4
-BIN_MAX = 4
-
-
-
-model = UNet().to(DEVICE)
-model.load_state_dict(torch.load('model.pth'))
-
-autoencoder = Autoencoder(input_dim=(1, 28, 28), latent_dim=(1, 8, 8)).to(DEVICE)
-autoencoder.load_state_dict(torch.load('autoencoder.pth'))
-
-
-
-
-# plt.figure()
-# plt.plot(BETA, label='beta')
-# plt.plot(ALPHA, label='alpha')
-# plt.plot(ALPHA_BAR, label='alpha_bar')
-# plt.legend()
-# plt.title('Alpha, Beta and Alpha_bar schedules')
-# plt.savefig('plots/alpha_beta.tmp.png')
-
-
-mnist = datasets.MNIST('data', train=True, download=True)
-img, label = mnist[np.random.randint(0, len(mnist))]
-img = np.array(img) / 255 * 2 - 1
-img = torch.tensor(img, device=DEVICE, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
-
-encoded = autoencoder.encode(img)
-
-img = encoded.squeeze().cpu().detach().numpy()
-
-plt.figure()
-plt.imshow(img, cmap='gray')
-plt.title('Image')
-plt.axis('off')
-plt.savefig('plots/img.tmp.png')
-
-
-plt.figure()
-plt.hist(img.flatten(), bins=NB_BINS, range=(BIN_MIN, BIN_MAX))
-plt.yscale('log')
-plt.title('Image histogram')
-plt.savefig('plots/img_hist.tmp.png')
-
-def norm_dist(x, mean, std):
- return np.exp(-0.5 * ((x - mean) / std) ** 2) / (std * np.sqrt(2 * np.pi))
-
-x_norm = np.linspace(BIN_MIN, BIN_MAX, 100)
-y_norm = norm_dist(x_norm, 0, 1) * 8**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:])
-
-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_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/{t:04d}.tmp.png')
-os.system('convert -delay 20 -loop 0 plots/diffusion/*.png plots/diffusion.tmp.gif')
-
-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)
-vec[0, label] = 1
-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 -loop 0 -reverse plots/diffusion_inverse/*.png plots/diffusion_inverse.tmp.gif')
-
-exit(0)
-
-tpause = {150: 'xt', 20: 'mu', 50: 'xt_1'}
-
-xT = xt = torch.randn(1, 1, 28, 28, device=DEVICE)
-vec = torch.zeros(1, 10).to(DEVICE)
-vec[0, label] = 1
-
-for t in 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 in tpause:
- plt.figure(figsize=(5, 5))
- plt.imshow(tensor_to_image(xt[0]), cmap='gray')
- plt.axis('off')
- plt.tight_layout()
- plt.savefig(f'plots/{tpause[t]}.tmp.png')
-
-plt.figure(figsize=(5, 5))
-plt.imshow(tensor_to_image(xT[0]), cmap='gray')
-plt.axis('off')
-plt.tight_layout()
-plt.savefig('plots/xT.tmp.png')
-
-plt.figure(figsize=(5, 5))
-plt.imshow(tensor_to_image(xt[0]), cmap='gray')
-plt.axis('off')
-plt.tight_layout()
-plt.savefig('plots/x0.tmp.png')