3 # Any copyright is dedicated to the Public Domain.
4 # https://creativecommons.org/publicdomain/zero/1.0/
6 # Written by Francois Fleuret <francois@fleuret.org>
8 # Minimal implementation of Jonathan Ho, Ajay Jain, Pieter Abbeel
9 # "Denoising Diffusion Probabilistic Models" (2020)
11 # https://arxiv.org/abs/2006.11239
13 import matplotlib.pyplot as plt
17 ######################################################################
21 result = torch.empty(nb).normal_(0, std)
22 result = result + torch.sign(torch.rand(result.size()) - p) / 2
25 ######################################################################
27 model = nn.Sequential(
35 ######################################################################
42 train_input = sample_phi(nb_samples)[:, None]
45 beta = torch.linspace(1e-4, 0.02, T)
47 alpha_bar = alpha.log().cumsum(0).exp()
50 for k in range(nb_epochs):
52 optimizer = torch.optim.Adam(model.parameters(), lr = 1e-4 * (1 - k / nb_epochs) )
54 for x0 in train_input.split(batch_size):
55 t = torch.randint(T, (x0.size(0), 1))
56 eps = torch.randn(x0.size())
57 input = alpha_bar[t].sqrt() * x0 + (1 - alpha_bar[t]).sqrt() * eps
58 input = torch.cat((input, 2 * t / T - 1), 1)
60 loss = (eps - output).pow(2).mean()
65 acc_loss += loss.item()
67 if k%10 == 0: print(k, loss.item())
69 ######################################################################
72 x = torch.randn(10000, 1)
74 for t in range(T-1, -1, -1):
75 z = torch.zeros(x.size()) if t == 0 else torch.randn(x.size())
76 input = torch.cat((x, torch.ones(x.size(0), 1) * 2 * t / T - 1), 1)
77 x = 1 / alpha[t].sqrt() * (x - (1 - alpha[t])/(1 - alpha_bar[t]).sqrt() * model(input)) + sigma[t] * z
79 ######################################################################
83 ax = fig.add_subplot(1, 1, 1)
84 ax.set_xlim(-1.25, 1.25)
86 d = train_input.flatten().detach().numpy()
87 ax.hist(d, 25, (-1, 1),
89 histtype = 'stepfilled', color = 'lightblue', label = 'Train')
91 d = x.flatten().detach().numpy()
92 ax.hist(d, 25, (-1, 1),
94 histtype = 'step', color = 'red', label = 'Synthesis')
96 ax.legend(frameon = False, loc = 2)
98 filename = 'diffusion.pdf'
99 print(f'saving {filename}')
100 fig.savefig(filename, bbox_inches='tight')
104 ######################################################################