From 11f27032ba4e5f97f8fc57e29de7791910e4a37f Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Sun, 6 Sep 2020 14:22:44 -0700 Subject: [PATCH] offer training class to easily train model off an image directory --- README.md | 33 +++++++ denoising_diffusion_pytorch/__init__.py | 2 +- .../denoising_diffusion_pytorch.py | 92 ++++++++++++++++++- setup.py | 4 +- 4 files changed, 125 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index b1035d9..56b7fbe 100644 --- a/README.md +++ b/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 diff --git a/denoising_diffusion_pytorch/__init__.py b/denoising_diffusion_pytorch/__init__.py index f78d79f..06ba76e 100644 --- a/denoising_diffusion_pytorch/__init__.py +++ b/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 diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 3b005c5..f7af30f 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/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') diff --git a/setup.py b/setup.py index 26044fa..14684ea 100644 --- a/setup.py +++ b/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=[