From: Francois Fleuret Date: Thu, 11 Aug 2022 20:44:23 +0000 (+0200) Subject: Initial commit X-Git-Url: https://fleuret.org/cgi-bin/gitweb/gitweb.cgi?a=commitdiff_plain;h=bc937c74ad8cbeede2c2ae97cda72eaa3e9bb4f3;p=pytorch.git Initial commit --- diff --git a/minidiffusion.py b/minidiffusion.py new file mode 100755 index 0000000..cbdb142 --- /dev/null +++ b/minidiffusion.py @@ -0,0 +1,89 @@ +#!/usr/bin/env python + +# Any copyright is dedicated to the Public Domain. +# https://creativecommons.org/publicdomain/zero/1.0/ + +# Written by Francois Fleuret + +import matplotlib.pyplot as plt +import torch +from torch import nn + +###################################################################### + +def sample_phi(nb): + p, std = 0.3, 0.2 + result = torch.empty(nb).normal_(0, std) + result = result + torch.sign(torch.rand(result.size()) - p) / 2 + return result + +###################################################################### + +model = nn.Sequential( + nn.Linear(2, 32), + nn.ReLU(), + nn.Linear(32, 32), + nn.ReLU(), + nn.Linear(32, 1), +) + +###################################################################### +# Train + +nb_samples = 25000 +nb_epochs = 250 +batch_size = 100 + +train_input = sample_phi(nb_samples)[:, None] + +T = 1000 +beta = torch.linspace(1e-4, 0.02, T) +alpha = 1 - beta +alpha_bar = alpha.log().cumsum(0).exp() +sigma = beta.sqrt() + +for k in range(nb_epochs): + acc_loss = 0 + optimizer = torch.optim.Adam(model.parameters(), lr = 1e-4 * (1 - k / nb_epochs) ) + + for x0 in train_input.split(batch_size): + t = torch.randint(T, (x0.size(0), 1)) + eps = torch.randn(x0.size()) + input = alpha_bar[t].sqrt() * x0 + (1 - alpha_bar[t]).sqrt() * eps + input = torch.cat((input, 2 * t / T - 1), 1) + output = model(input) + loss = (eps - output).pow(2).mean() + optimizer.zero_grad() + loss.backward() + optimizer.step() + + acc_loss += loss.item() + + if k%10 == 0: print(k, loss.item()) + +###################################################################### +# Plot + +x = torch.randn(10000, 1) + +for t in range(T-1, -1, -1): + z = torch.zeros(x.size()) if t == 0 else torch.randn(x.size()) + input = torch.cat((x, torch.ones(x.size(0), 1) * 2 * t / T - 1), 1) + x = 1 / alpha[t].sqrt() * (x - (1 - alpha[t])/(1 - alpha_bar[t]).sqrt() * model(input)) + sigma[t] * z + +fig = plt.figure() +ax = fig.add_subplot(1, 1, 1) +ax.set_xlim(-1.25, 1.25) + +d = train_input.flatten().detach().numpy() +ax.hist(d, 25, (-1, 1), histtype = 'stepfilled', color = 'lightblue', density = True, label = 'Train') + +d = x.flatten().detach().numpy() +ax.hist(d, 25, (-1, 1), histtype = 'step', color = 'red', density = True, label = 'Synthesis') + +filename = 'diffusion.pdf' +fig.savefig(filename, bbox_inches='tight') + +plt.show() + +######################################################################