diff options
Diffstat (limited to 'plots.py')
| -rw-r--r-- | plots.py | 174 |
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') |
