From 22ab4eba0570fa4a2463835c69037aa13838dc4a Mon Sep 17 00:00:00 2001 From: gdamms Date: Fri, 3 May 2024 17:58:48 +0200 Subject: use my understand not paper formula --- main.py | 60 +++++++++++++++++++++++++++++------------------------------- 1 file changed, 29 insertions(+), 31 deletions(-) (limited to 'main.py') diff --git a/main.py b/main.py index f63bc6a..6f368bf 100644 --- a/main.py +++ b/main.py @@ -78,23 +78,16 @@ class UNet(nn.Module): 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) +DIFFU_STEPS = 20 +BETA = torch.linspace(1e-2, 3e-1, DIFFU_STEPS, device=DEVICE) BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA)) ALPHA = 1 - BETA ALPHA_BAR = torch.cumprod(ALPHA, dim=0) -ALPHA_BAR_BAR = torch.cumsum(torch.sqrt(ALPHA_BAR), dim=0) - -plt.figure() -plt.plot(BETA.cpu().detach().numpy(), label="beta") -plt.plot(ALPHA.cpu().detach().numpy(), label="alpha") -plt.plot(ALPHA_BAR.cpu().detach().numpy(), label="alpha_bar") -# plt.plot(ALPHA_BAR_BAR.cpu().detach().numpy(), label="aplha_bar_bar") -plt.plot(ALPHA_BAR_BAR.cpu().detach().numpy() * (1 - ALPHA).cpu().detach().numpy(), label="sigma_theta") -plt.legend() -plt.xlim(0, DIFFU_STEPS) -# plt.ylim(0, 1) -plt.savefig("beta_alpha.tmp.png") +CUM_SQ_SUM = [] +for t in range(len(ALPHA)): + cum_sq_sum = sum([torch.prod(ALPHA[s+2:t+1]) * (1 - ALPHA[s+1])**2 for s in range(t)]) + CUM_SQ_SUM.append(cum_sq_sum) +CUM_SQ_SUM = torch.tensor(CUM_SQ_SUM, device=DEVICE) def q_xt_xt_1(xt_1, t): @@ -105,19 +98,26 @@ def q_xt_xt_1(xt_1, t): std = 1 - alpha xt = torch.distributions.Normal(mean, std).sample() - print('xt_1 xt', xt_1.min(), xt_1.max(), xt.min(), xt.max()) - return xt +# def q_xt_x0(x0, t): +# t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t + +# alpha = ALPHA[t_ind] +# alpha_bar = ALPHA_BAR[t_ind] +# alpha_bar_bar = ALPHA_BAR_BAR[t_ind - 1] +# mean = torch.sqrt(alpha_bar) * x0 +# std = alpha_bar_bar * (1 - alpha) +# return torch.distributions.Normal(mean, std).sample() def q_xt_x0(x0, t): t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t - alpha = ALPHA[t_ind] alpha_bar = ALPHA_BAR[t_ind] - alpha_bar_bar = ALPHA_BAR_BAR[t_ind - 1] + cum_sq_sum = CUM_SQ_SUM[t_ind] + mean = torch.sqrt(alpha_bar) * x0 - std = alpha_bar_bar * (1 - alpha) + std = torch.sqrt(cum_sq_sum) return torch.distributions.Normal(mean, std).sample() @@ -190,7 +190,7 @@ def loss(y_pred, y_true): def forward_diffusion(x0): x = x0.clone()[0] xs = [x.cpu().detach().numpy()] - for t in range(1, DIFFU_STEPS): + for t in range(1, DIFFU_STEPS+1): x = q_xt_xt_1(x, t) xs.append(x.cpu().detach().numpy()) return xs @@ -222,9 +222,8 @@ if __name__ == '__main__': ax.imshow(x, cmap="gray") ax.axis("off") - for t in range(1, DIFFU_STEPS): + for t in range(1, DIFFU_STEPS+1): x = q_xt_x0(img, t)[0].cpu() - # print('x0', x.min(), x.max()) if t not in plots_id: continue plot_i = plots_id.index(t) @@ -233,8 +232,6 @@ if __name__ == '__main__': ax.axis("off") fig.tight_layout() fig.savefig("img.tmp.png") - exit() - ############ # Training # @@ -251,22 +248,23 @@ if __name__ == '__main__': # Define the trainiself.encodevec = nn.Linear(10, 28*28)ng dataset. train_dataset = MNISTDiffusionDataset(train=True) - train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4) + train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, + num_workers=4, persistent_workers=True) trainer = Trainer() criterion = loss - epochs = 30 + epochs = 1 - # # Train the model. - # trainer.train(model, train_loader, epochs, optimizer, criterion) + # Train the model. + trainer.train(model, train_loader, epochs, optimizer, criterion) - # # Save the model. - # torch.save(model.state_dict(), "model.pth") + # Save the model. + torch.save(model.state_dict(), "model.pth") ############## # Evaluation # ############## - nb_plots = 5 + nb_plots = 10 ti_plots = np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int) n_values = [i for i in range(10)] -- cgit v1.3.1