offer training class to easily train model off an image directory

README.md

@@ -38,6 +38,39 @@ sampled_images = diffusion.p_sample_loop((1, 3, 128, 128))
 sampled_images.shape # (1, 3, 128, 128)
 ```
 
+Or, if you simply want to pass in a folder name and the desired image dimensions, you can use the `Trainer` class to easily train a model.
+
+```python
+from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer
+
+model = Unet(
+    dim = 64,
+    dim_mults = (1, 2, 4, 8)
+).cuda()
+
+diffusion = GaussianDiffusion(
+    model,
+    beta_start = 0.0001,
+    beta_end = 0.02,
+    num_diffusion_timesteps = 1000,   # number of steps
+    loss_type = 'l1'                  # L1 or L2
+).cuda()
+
+trainer = Trainer(
+    diffusion,
+    'path/to/your/images',
+    image_size = 128,
+    train_batch_size = 32,
+    train_lr = 3e-4,
+    train_num_steps = 100000,
+    gradient_accumulate_every = 1
+)
+
+trainer.train()
+```
+
+Todo: Command line tool for one-line training
+
 ## Citations
 
 ```bibtex

denoising_diffusion_pytorch/__init__.py

@@ -1 +1 @@
-from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, Unet
+from denoising_diffusion_pytorch.denoising_diffusion_pytorch import GaussianDiffusion, Unet, Trainer

denoising_diffusion_pytorch/denoising_diffusion_pytorch.py

@@ -1,14 +1,25 @@
 import math
 import torch
-from inspect import isfunction
-from functools import partial
 from torch import nn, einsum
 import torch.nn.functional as F
+from inspect import isfunction
+from functools import partial
+
+from torch.utils import data
+from pathlib import Path
+from torch.optim import Adam
+from torchvision import transforms, utils
+from PIL import Image
 
 import numpy as np
 from tqdm import tqdm
 from einops import rearrange
 
+# constants
+
+SAVE_AND_SAMPLE_EVERY = 1000
+EXTS = ['jpg', 'png']
+
 # helpers functions
 
 def exists(x):
@@ -19,8 +30,10 @@ def default(val, d):
         return val
     return d() if isfunction(d) else d
 
-def normal_kl(mean1, logvar1, mean2, logvar2):
-    return 0.5 * (-1. + logvar2 - logvar1 + torch.exp(logvar1 - logvar2) + torch.exp(-logvar2) * (mean1 - mean2) ** 2)
+def cycle(dl):
+    while True:
+        for data in dl:
+            yield data
 
 # small helper modules
 
@@ -324,3 +337,74 @@ class GaussianDiffusion(nn.Module):
         b, *_, device = *x.shape, x.device
         t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
         return self.p_losses(x, t, *args, **kwargs)
+
+# dataset classes
+
+class Dataset(data.Dataset):
+    def __init__(self, folder, image_size):
+        super().__init__()
+        self.folder = folder
+        self.image_size = image_size
+        self.paths = [p for ext in EXTS for p in Path(f'{folder}').glob(f'**/*.{ext}')]
+
+        self.transform = transforms.Compose([
+            transforms.Resize(image_size),
+            transforms.RandomHorizontalFlip(),
+            transforms.CenterCrop(image_size),
+            transforms.ToTensor()
+        ])
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, index):
+        path = self.paths[index]
+        img = Image.open(path)
+        return self.transform(img)
+
+# trainer class
+
+class Trainer(object):
+    def __init__(
+        self,
+        diffusion_model,
+        folder,
+        *,
+        image_size = 128,
+        train_batch_size = 32,
+        train_lr = 3e-4,
+        train_num_steps = 100000,
+        gradient_accumulate_every = 1
+    ):
+        super().__init__()
+        self.model = diffusion_model
+        self.image_size = image_size
+        self.gradient_accumulate_every = gradient_accumulate_every
+        self.train_num_steps = train_num_steps
+
+        self.ds = Dataset(folder, image_size)
+        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)
+
+    def train(self):
+        ind = 0
+
+        while ind < self.train_num_steps:
+            for i in range(self.gradient_accumulate_every):
+                data = next(self.dl).cuda()
+                loss = self.model(data)
+                print(f'{ind}: {loss.item()}')
+                loss.backward()
+
+            self.opt.step()
+            self.opt.zero_grad()
+
+            if ind % SAVE_AND_SAMPLE_EVERY == 0:
+                milestone = ind // SAVE_AND_SAMPLE_EVERY
+                all_images = self.model.p_sample_loop((64, 3, self.image_size, self.image_size))
+                utils.save_image(all_images, f'./sample-{milestone}.png', nrow=8)
+                torch.save(model.state_dict(), f'./model-{milestone}.pt')
+
+            ind += 1
+
+        print('training completed')

setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'denoising-diffusion-pytorch',
   packages = find_packages(),
-  version = '0.0.2',
+  version = '0.1.0',
   license='MIT',
   description = 'Denoising Diffusion Probabilistic Models - Pytorch',
   author = 'Phil Wang',
@@ -16,7 +16,9 @@ setup(
   install_requires=[
     'einops',
     'numpy',
+    'pillow',
     'torch',
+    'torchvision',
     'tqdm'
   ],
   classifiers=[