add gaussian diffusion where model predicts both noise and x_start, with a learned weighting between the two (experimental)

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -603,37 +603,34 @@ class Trainer(object):
         self.scaler.load_state_dict(data['scaler'])
 
     def train(self):
-        with tqdm(initial = self.step, total = self.train_num_steps) as pbar:
+        while self.step < self.train_num_steps:
+            for i in range(self.gradient_accumulate_every):
+                data = next(self.dl).cuda()
 
-            while self.step < self.train_num_steps:
-                for i in range(self.gradient_accumulate_every):
-                    data = next(self.dl).cuda()
+                with autocast(enabled = self.amp):
+                    loss = self.model(data)
+                    self.scaler.scale(loss / self.gradient_accumulate_every).backward()
 
-                    with autocast(enabled = self.amp):
-                        loss = self.model(data)
-                        self.scaler.scale(loss / self.gradient_accumulate_every).backward()
+                print(f'{self.step}: {loss.item()}')
 
-                    pbar.set_description(f'loss: {loss.item():.4f}')
+            self.scaler.step(self.opt)
+            self.scaler.update()
+            self.opt.zero_grad()
 
-                self.scaler.step(self.opt)
-                self.scaler.update()
-                self.opt.zero_grad()
+            if self.step % self.update_ema_every == 0:
+                self.step_ema()
 
-                if self.step % self.update_ema_every == 0:
-                    self.step_ema()
+            if self.step != 0 and self.step % self.save_and_sample_every == 0:
+                self.ema_model.eval()
 
-                if self.step != 0 and self.step % self.save_and_sample_every == 0:
-                    self.ema_model.eval()
+                milestone = self.step // self.save_and_sample_every
+                batches = num_to_groups(36, self.batch_size)
+                all_images_list = list(map(lambda n: self.ema_model.sample(batch_size=n), batches))
+                all_images = torch.cat(all_images_list, dim=0)
+                all_images = unnormalize_to_zero_to_one(all_images)
+                utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = 6)
+                self.save(milestone)
 
-                    milestone = self.step // self.save_and_sample_every
-                    batches = num_to_groups(36, self.batch_size)
-                    all_images_list = list(map(lambda n: self.ema_model.sample(batch_size=n), batches))
-                    all_images = torch.cat(all_images_list, dim=0)
-                    all_images = unnormalize_to_zero_to_one(all_images)
-                    utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = 6)
-                    self.save(milestone)
+            self.step += 1
 
-                self.step += 1
-                pbar.update(1)
-
-        print('training complete')
+        print('training completed')

denoising_diffusion_pytorch/weighted_objective_gaussian_diffusion.py

@@ -3,7 +3,7 @@ from inspect import isfunction
 from torch import nn, einsum
 from einops import rearrange
 
-from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion
+from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, extract, unnormalize_to_zero_to_one
 
 # helper functions
 

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.15.3',
+  version = '0.15.1',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',