From e8500a3f0cec4be59442e2b3bdbe692b04a9524a Mon Sep 17 00:00:00 2001 From: Francois Fleuret Date: Tue, 16 Aug 2022 22:51:29 +0200 Subject: [PATCH] Update. --- minidiffusion.py | 105 ++++++++++++++++++++++++++++++----------------- 1 file changed, 67 insertions(+), 38 deletions(-) diff --git a/minidiffusion.py b/minidiffusion.py index 27842d9..7327522 100755 --- a/minidiffusion.py +++ b/minidiffusion.py @@ -11,6 +11,7 @@ import matplotlib.pyplot as plt import torch, torchvision from torch import nn +from torch.nn import functional as F device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') @@ -113,6 +114,9 @@ parser.add_argument('--data', type = str, default = 'gaussian_mixture', help = f'Toy data-set to use: {data_list}') +parser.add_argument('--no_window', + action='store_true', default = False) + args = parser.parse_args() if args.seed >= 0: @@ -223,7 +227,7 @@ print(f'nb_parameters {sum([ p.numel() for p in model.parameters() ])}') ###################################################################### # Generate -def generate(size, alpha, alpha_bar, sigma, model, train_mean, train_std): +def generate(size, T, alpha, alpha_bar, sigma, model, train_mean, train_std): with torch.no_grad(): @@ -280,72 +284,97 @@ if ema is not None: ema.copy_to_model() model.eval() -if train_input.dim() == 2: +######################################## +# Nx1 -> histogram +if train_input.dim() == 2 and train_input.size(1) == 1: fig = plt.figure() ax = fig.add_subplot(1, 1, 1) - # Nx1 -> histogram - if train_input.size(1) == 1: + x = generate((10000, 1), T, alpha, alpha_bar, sigma, + model, train_mean, train_std) - x = generate((10000, 1), alpha, alpha_bar, sigma, - model, train_mean, train_std) + ax.set_xlim(-1.25, 1.25) + ax.spines.right.set_visible(False) + ax.spines.top.set_visible(False) - ax.set_xlim(-1.25, 1.25) - ax.spines.right.set_visible(False) - ax.spines.top.set_visible(False) + d = train_input.flatten().detach().to('cpu').numpy() + ax.hist(d, 25, (-1, 1), + density = True, + histtype = 'bar', edgecolor = 'white', color = 'lightblue', label = 'Train') - d = train_input.flatten().detach().to('cpu').numpy() - ax.hist(d, 25, (-1, 1), - density = True, - histtype = 'stepfilled', color = 'lightblue', label = 'Train') + d = x.flatten().detach().to('cpu').numpy() + ax.hist(d, 25, (-1, 1), + density = True, + histtype = 'step', color = 'red', label = 'Synthesis') - d = x.flatten().detach().to('cpu').numpy() - ax.hist(d, 25, (-1, 1), - density = True, - histtype = 'step', color = 'red', label = 'Synthesis') + ax.legend(frameon = False, loc = 2) - ax.legend(frameon = False, loc = 2) + filename = f'diffusion_{args.data}.pdf' + print(f'saving {filename}') + fig.savefig(filename, bbox_inches='tight') - # Nx2 -> scatter plot - elif train_input.size(1) == 2: + if not args.no_window and hasattr(plt.get_current_fig_manager(), 'window'): + plt.get_current_fig_manager().window.setGeometry(2, 2, 2048, 768) + plt.show() - x = generate((1000, 2), alpha, alpha_bar, sigma, - model, train_mean, train_std) +######################################## +# Nx2 -> scatter plot +elif train_input.dim() == 2 and train_input.size(1) == 2: - ax.set_xlim(-1.5, 1.5) - ax.set_ylim(-1.5, 1.5) - ax.set(aspect = 1) - ax.spines.right.set_visible(False) - ax.spines.top.set_visible(False) + fig = plt.figure() + ax = fig.add_subplot(1, 1, 1) - d = x.detach().to('cpu').numpy() - ax.scatter(d[:, 0], d[:, 1], - s = 2.0, color = 'red', label = 'Synthesis') + x = generate((1000, 2), T, alpha, alpha_bar, sigma, + model, train_mean, train_std) + + ax.set_xlim(-1.5, 1.5) + ax.set_ylim(-1.5, 1.5) + ax.set(aspect = 1) + ax.spines.right.set_visible(False) + ax.spines.top.set_visible(False) - d = train_input[:x.size(0)].detach().to('cpu').numpy() - ax.scatter(d[:, 0], d[:, 1], - s = 2.0, color = 'gray', label = 'Train') + d = train_input[:x.size(0)].detach().to('cpu').numpy() + ax.scatter(d[:, 0], d[:, 1], + s = 2.5, color = 'gray', label = 'Train') - ax.legend(frameon = False, loc = 2) + d = x.detach().to('cpu').numpy() + ax.scatter(d[:, 0], d[:, 1], + s = 2.0, color = 'red', label = 'Synthesis') + + ax.legend(frameon = False, loc = 2) filename = f'diffusion_{args.data}.pdf' print(f'saving {filename}') fig.savefig(filename, bbox_inches='tight') - if hasattr(plt.get_current_fig_manager(), 'window'): + if not args.no_window and hasattr(plt.get_current_fig_manager(), 'window'): plt.get_current_fig_manager().window.setGeometry(2, 2, 1024, 768) plt.show() +######################################## # NxCxHxW -> image elif train_input.dim() == 4: - x = generate((128,) + train_input.size()[1:], alpha, alpha_bar, sigma, + x = generate((128,) + train_input.size()[1:], T, alpha, alpha_bar, sigma, model, train_mean, train_std) - x = 1 - x.clamp(min = 0, max = 255) / 255 + + x = torchvision.utils.make_grid(x.clamp(min = 0, max = 255), + nrow = 16, padding = 1, pad_value = 64) + x = F.pad(x, pad = (2, 2, 2, 2), value = 64)[None] + + t = torchvision.utils.make_grid(train_input[:128], + nrow = 16, padding = 1, pad_value = 64) + t = F.pad(t, pad = (2, 2, 2, 2), value = 64)[None] + + result = 1 - torch.cat((t, x), 2) / 255 filename = f'diffusion_{args.data}.png' print(f'saving {filename}') - torchvision.utils.save_image(x, filename, nrow = 16, pad_value = 0.8) + torchvision.utils.save_image(result, filename) + +else: + + print(f'cannot plot result of size {train_input.size()}') ###################################################################### -- 2.39.5