Initial commit

diff --git a/minidiffusion.py b/minidiffusion.py
new file mode 100755 (executable)
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 <francois@fleuret.org>
+
+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()
+
+######################################################################