From be287add0311cc66345e5a26e297b4fe30310398 Mon Sep 17 00:00:00 2001
From: Francois Fleuret <francois@fleuret.org>
Date: Fri, 12 Aug 2022 23:05:22 +0200
Subject: [PATCH] Update.

---
 minidiffusion.py | 110 ++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 85 insertions(+), 25 deletions(-)

diff --git a/minidiffusion.py b/minidiffusion.py
index a386a12..0f5948e 100755
--- a/minidiffusion.py
+++ b/minidiffusion.py
@@ -10,36 +10,70 @@
 #
 # https://arxiv.org/abs/2006.11239
 
+import math
 import matplotlib.pyplot as plt
 import torch
 from torch import nn
 
 ######################################################################
 
-def sample_phi(nb):
+class EMA:
+    def __init__(self, model, decay = 0.9999):
+        self.model = model
+        self.decay = decay
+        self.ema = { }
+        with torch.no_grad():
+            for p in model.parameters():
+                self.ema[p] = p.clone()
+
+    def step(self):
+        with torch.no_grad():
+            for p in self.model.parameters():
+                self.ema[p].copy_(self.decay * self.ema[p] + (1 - self.decay) * p)
+
+    def copy(self):
+        with torch.no_grad():
+            for p in self.model.parameters():
+                p.copy_(self.ema[p])
+
+######################################################################
+
+def sample_gaussian_mixture(nb):
     p, std = 0.3, 0.2
-    result = torch.empty(nb).normal_(0, std)
+    result = torch.empty(nb, 1).normal_(0, std)
     result = result + torch.sign(torch.rand(result.size()) - p) / 2
     return result
 
+def sample_arc(nb):
+    theta = torch.rand(nb) * math.pi
+    rho = torch.rand(nb) * 0.1 + 0.7
+    result = torch.empty(nb, 2)
+    result[:, 0] = theta.cos() * rho
+    result[:, 1] = theta.sin() * rho
+    return result
+
+######################################################################
+# Train
+
+nb_samples = 25000
+
+train_input = sample_gaussian_mixture(nb_samples)
+#train_input = sample_arc(nb_samples)
+
 ######################################################################
 
+nh = 64
+
 model = nn.Sequential(
-    nn.Linear(2, 32),
+    nn.Linear(train_input.size(1) + 1, nh),
     nn.ReLU(),
-    nn.Linear(32, 32),
+    nn.Linear(nh, nh),
     nn.ReLU(),
-    nn.Linear(32, 1),
+    nn.Linear(nh, train_input.size(1)),
 )
 
-######################################################################
-# Train
-
-nb_samples = 25000
-nb_epochs = 250
-batch_size = 100
-
-train_input = sample_phi(nb_samples)[:, None]
+nb_epochs = 50
+batch_size = 25
 
 T = 1000
 beta = torch.linspace(1e-4, 0.02, T)
@@ -47,10 +81,12 @@ alpha = 1 - beta
 alpha_bar = alpha.log().cumsum(0).exp()
 sigma = beta.sqrt()
 
+ema = EMA(model)
+
 for k in range(nb_epochs):
 
     acc_loss = 0
-    optimizer = torch.optim.Adam(model.parameters(), lr = 1e-4 * (1 - k / nb_epochs) )
+    optimizer = torch.optim.Adam(model.parameters(), lr = 1e-3)
 
     for x0 in train_input.split(batch_size):
         t = torch.randint(T, (x0.size(0), 1))
@@ -65,12 +101,16 @@ for k in range(nb_epochs):
 
         acc_loss += loss.item()
 
+        ema.step()
+
     if k%10 == 0: print(k, loss.item())
 
+ema.copy()
+
 ######################################################################
 # Generate
 
-x = torch.randn(10000, 1)
+x = torch.randn(10000, train_input.size(1))
 
 for t in range(T-1, -1, -1):
     z = torch.zeros(x.size()) if t == 0 else torch.randn(x.size())
@@ -83,24 +123,44 @@ for t in range(T-1, -1, -1):
 
 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),
-        density = True,
-        histtype = 'stepfilled', color = 'lightblue', label = 'Train')
+if train_input.size(1) == 1:
+
+    ax.set_xlim(-1.25, 1.25)
+
+    d = train_input.flatten().detach().numpy()
+    ax.hist(d, 25, (-1, 1),
+            density = True,
+            histtype = 'stepfilled', color = 'lightblue', label = 'Train')
+
+    d = x.flatten().detach().numpy()
+    ax.hist(d, 25, (-1, 1),
+            density = True,
+            histtype = 'step', color = 'red', label = 'Synthesis')
+
+    ax.legend(frameon = False, loc = 2)
+
+elif train_input.size(1) == 2:
+
+    ax.set_xlim(-1.25, 1.25)
+    ax.set_ylim(-1.25, 1.25)
+    ax.set(aspect = 1)
+
+    d = train_input[:200].detach().numpy()
+    ax.scatter(d[:, 0], d[:, 1],
+               color = 'lightblue', label = 'Train')
 
-d = x.flatten().detach().numpy()
-ax.hist(d, 25, (-1, 1),
-        density = True,
-        histtype = 'step', color = 'red', label = 'Synthesis')
+    d = x[:200].detach().numpy()
+    ax.scatter(d[:, 0], d[:, 1],
+               color = 'red', label = 'Synthesis')
 
-ax.legend(frameon = False, loc = 2)
+    ax.legend(frameon = False, loc = 2)
 
 filename = 'diffusion.pdf'
 print(f'saving {filename}')
 fig.savefig(filename, bbox_inches='tight')
 
+plt.get_current_fig_manager().window.setGeometry(2, 2, 1024, 768)
 plt.show()
 
 ######################################################################
-- 
2.39.5