+from torch.nn import functional as F
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+print(f'device {device}')
+
+######################################################################
+
+def sample_gaussian_mixture(nb):
+ p, std = 0.3, 0.2
+ result = torch.randn(nb, 1) * std
+ result = result + torch.sign(torch.rand(result.size()) - p) / 2
+ return result
+
+def sample_ramp(nb):
+ result = torch.min(torch.rand(nb, 1), torch.rand(nb, 1))
+ return result
+
+def sample_two_discs(nb):
+ a = torch.rand(nb) * math.pi * 2
+ b = torch.rand(nb).sqrt()
+ q = (torch.rand(nb) <= 0.5).long()
+ b = b * (0.3 + 0.2 * q)
+ result = torch.empty(nb, 2)
+ result[:, 0] = a.cos() * b - 0.5 + q
+ result[:, 1] = a.sin() * b - 0.5 + q
+ return result
+
+def sample_disc_grid(nb):
+ a = torch.rand(nb) * math.pi * 2
+ b = torch.rand(nb).sqrt()
+ N = 4
+ q = (torch.randint(N, (nb,)) - (N - 1) / 2) / ((N - 1) / 2)
+ r = (torch.randint(N, (nb,)) - (N - 1) / 2) / ((N - 1) / 2)
+ b = b * 0.1
+ result = torch.empty(nb, 2)
+ result[:, 0] = a.cos() * b + q
+ result[:, 1] = a.sin() * b + r
+ return result
+
+def sample_spiral(nb):
+ u = torch.rand(nb)
+ rho = u * 0.65 + 0.25 + torch.rand(nb) * 0.15
+ theta = u * math.pi * 3
+ result = torch.empty(nb, 2)
+ result[:, 0] = theta.cos() * rho
+ result[:, 1] = theta.sin() * rho
+ return result
+
+def sample_mnist(nb):
+ train_set = torchvision.datasets.MNIST(root = './data/', train = True, download = True)
+ result = train_set.data[:nb].to(device).view(-1, 1, 28, 28).float()
+ return result
+
+samplers = {
+ f.__name__.removeprefix('sample_') : f for f in [
+ sample_gaussian_mixture,
+ sample_ramp,
+ sample_two_discs,
+ sample_disc_grid,
+ sample_spiral,
+ sample_mnist,
+ ]
+}
+
+######################################################################
+
+parser = argparse.ArgumentParser(
+ description = '''A minimal implementation of Jonathan Ho, Ajay Jain, Pieter Abbeel
+"Denoising Diffusion Probabilistic Models" (2020)
+https://arxiv.org/abs/2006.11239''',
+
+ formatter_class = argparse.ArgumentDefaultsHelpFormatter
+)
+
+parser.add_argument('--seed',
+ type = int, default = 0,
+ help = 'Random seed, < 0 is no seeding')
+
+parser.add_argument('--nb_epochs',
+ type = int, default = 100,
+ help = 'How many epochs')
+
+parser.add_argument('--batch_size',
+ type = int, default = 25,
+ help = 'Batch size')
+
+parser.add_argument('--nb_samples',
+ type = int, default = 25000,
+ help = 'Number of training examples')
+
+parser.add_argument('--learning_rate',
+ type = float, default = 1e-3,
+ help = 'Learning rate')
+
+parser.add_argument('--ema_decay',
+ type = float, default = 0.9999,
+ help = 'EMA decay, <= 0 is no EMA')
+
+data_list = ', '.join( [ str(k) for k in samplers ])
+
+parser.add_argument('--data',
+ type = str, default = 'gaussian_mixture',
+ help = f'Toy data-set to use: {data_list}')
+
+parser.add_argument('--no_window',
+ action='store_true', default = False)
+
+args = parser.parse_args()
+
+if args.seed >= 0:
+ # torch.backends.cudnn.deterministic = True
+ # torch.backends.cudnn.benchmark = False
+ # torch.use_deterministic_algorithms(True)
+ torch.manual_seed(args.seed)
+ if torch.cuda.is_available():
+ torch.cuda.manual_seed_all(args.seed)