for https://github.com/lucidrains/denoising-diffusion-pytorch/issues/36

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
+from einops import rearrange, repeat, reduce
 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 = losses.mean(dim = tuple(range(1, losses.ndim)))
+        losses = reduce(losses, 'b ... -> b', 'mean')
 
         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,6 +7,7 @@ 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
@@ -598,7 +599,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))
+        self.dl = cycle(data.DataLoader(self.ds, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count()))
         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.1',
+  version = '0.18.3',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',