allow for turning off horizontal flip augmentation

README.md

@@ -6,8 +6,6 @@ Implementation of <a href="https://arxiv.org/abs/2006.11239">Denoising Diffusion
 
 This implementation was transcribed from the official Tensorflow version <a href="https://github.com/hojonathanho/diffusion">here</a>
 
-Youtube AI Educators - <a href="https://www.youtube.com/watch?v=W-O7AZNzbzQ">Yannic Kilcher</a> | <a href="https://www.youtube.com/watch?v=344w5h24-h8">AI Coffeebreak with Letitia</a> | <a href="https://www.youtube.com/watch?v=HoKDTa5jHvg">Outlier</a>
-
 <a href="https://huggingface.co/blog/annotated-diffusion">Annotated code</a> by Research Scientists / Engineers from <a href="https://huggingface.co/">🤗 Huggingface</a>
 
 <img src="./sample.png" width="500px"><img>

denoising_diffusion_pytorch/continuous_time_gaussian_diffusion.py

@@ -5,7 +5,7 @@ import torch.nn.functional as F
 from torch.special import expm1
 
 from tqdm import tqdm
-from einops import rearrange, repeat, reduce
+from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 
 # helpers
@@ -268,7 +268,7 @@ class ContinuousTimeGaussianDiffusion(nn.Module):
         model_out = self.denoise_fn(x, log_snr)
 
         losses = self.loss_fn(model_out, noise, reduction = 'none')
-        losses = reduce(losses, 'b ... -> b', 'mean')
+        losses = losses.mean(dim = tuple(range(1, losses.ndim)))
 
         if self.p2_loss_weight_gamma >= 0:
             # following eq 8. in https://arxiv.org/abs/2204.00227

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -7,7 +7,6 @@ from inspect import isfunction
 from functools import partial
 
 from torch.utils import data
-from multiprocessing import cpu_count
 from torch.cuda.amp import autocast, GradScaler
 
 from pathlib import Path
@@ -16,7 +15,7 @@ from torchvision import transforms, utils
 from PIL import Image
 
 from tqdm import tqdm
-from einops import rearrange, reduce
+from einops import rearrange
 from einops.layers.torch import Rearrange
 
 # helpers functions
@@ -367,9 +366,7 @@ class GaussianDiffusion(nn.Module):
         timesteps = 1000,
         loss_type = 'l1',
         objective = 'pred_noise',
-        beta_schedule = 'cosine',
-        p2_loss_weight_gamma = 0., # p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time - 1. is recommended
-        p2_loss_weight_k = 1
+        beta_schedule = 'cosine'
     ):
         super().__init__()
         assert not (type(self) == GaussianDiffusion and denoise_fn.channels != denoise_fn.out_dim)
@@ -424,10 +421,6 @@ class GaussianDiffusion(nn.Module):
         register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
         register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod))
 
-        # calculate p2 reweighting
-
-        register_buffer('p2_loss_weight', (p2_loss_weight_k + alphas_cumprod / (1 - alphas_cumprod)) ** -p2_loss_weight_gamma)
-
     def predict_start_from_noise(self, x_t, t, noise):
         return (
             extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t -
@@ -534,11 +527,8 @@ class GaussianDiffusion(nn.Module):
         else:
             raise ValueError(f'unknown objective {self.objective}')
 
-        loss = self.loss_fn(model_out, target, reduction = 'none')
-        loss = reduce(loss, 'b ... -> b (...)', 'mean')
-
-        loss = loss * extract(self.p2_loss_weight, t, loss.shape)
-        return loss.mean()
+        loss = self.loss_fn(model_out, target)
+        return loss
 
     def forward(self, img, *args, **kwargs):
         b, c, h, w, device, img_size, = *img.shape, img.device, self.image_size
@@ -608,7 +598,7 @@ class Trainer(object):
         self.train_num_steps = train_num_steps
 
         self.ds = Dataset(folder, image_size, augment_horizontal_flip = augment_horizontal_flip)
-        self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count()))
+        self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle=True, pin_memory=True))
         self.opt = Adam(diffusion_model.parameters(), lr=train_lr)
 
         self.step = 0

setup.py

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