Merge pull request #204 from klasocki/main

Add Trainer1D for 1D diffusion

Author: lucidrains

README.md

@@ -106,11 +106,10 @@ $ accelerate launch train.py
 
 ### 1D Sequence
 
-By popular request, a 1D Unet + Gaussian Diffusion implementation. You will have to do the training code yourself
-
+By popular request, a 1D Unet + Gaussian Diffusion implementation.
 ```python
 import torch
-from denoising_diffusion_pytorch import Unet1D, GaussianDiffusion1D
+from denoising_diffusion_pytorch import Unet1D, GaussianDiffusion1D, Trainer1D
 
 model = Unet1D(
     dim = 64,
@@ -125,16 +124,33 @@ diffusion = GaussianDiffusion1D(
     objective = 'pred_v'
 )
 
-training_seq = torch.rand(8, 32, 128) # features are normalized from 0 to 1
+training_seq = torch.rand(64, 32, 128) # features are normalized from 0 to 1
 loss = diffusion(training_seq)
 loss.backward()
 
+# Or using trainer
+
+trainer = Trainer1D(
+    diffusion,
+    dataset = training_seq,
+    train_batch_size = 32,
+    train_lr = 8e-5,
+    train_num_steps = 700000,         # total training steps
+    gradient_accumulate_every = 2,    # gradient accumulation steps
+    ema_decay = 0.995,                # exponential moving average decay
+    amp = True,                       # turn on mixed precision
+)
+trainer.train()
+
 # after a lot of training
 
 sampled_seq = diffusion.sample(batch_size = 4)
 sampled_seq.shape # (4, 32, 128)
-```
 
+```
+`Trainer1D` does not evaluate the generated samples in any way since the type of data is not known. 
+You could consider adding a suitable metric to the training loop yourself after doing an editable install of this package
+`pip install -e .`.
 ## Citations
 
 ```bibtex

denoising_diffusion_pytorch/__init__.py

@@ -6,5 +6,5 @@
 from denoising_diffusion_pytorch.elucidated_diffusion import ElucidatedDiffusion
 from denoising_diffusion_pytorch.v_param_continuous_time_gaussian_diffusion import VParamContinuousTimeGaussianDiffusion
 
-from denoising_diffusion_pytorch.denoising_diffusion_pytorch_1d import GaussianDiffusion1D, Unet1D
+from denoising_diffusion_pytorch.denoising_diffusion_pytorch_1d import GaussianDiffusion1D, Unet1D, Trainer1D
 

denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py

@@ -1,17 +1,25 @@
 import math
+from multiprocessing import cpu_count
+from pathlib import Path
 from random import random
 from functools import partial
 from collections import namedtuple
 
 import torch
+from accelerate import Accelerator
+from ema_pytorch import EMA
 from torch import nn, einsum
 import torch.nn.functional as F
 
 from einops import rearrange, reduce
 from einops.layers.torch import Rearrange
+from torch.optim import Adam
+from torch.utils.data import Dataset, DataLoader
 
 from tqdm.auto import tqdm
 
+from denoising_diffusion_pytorch.version import __version__
+
 # constants
 
 ModelPrediction =  namedtuple('ModelPrediction', ['pred_noise', 'pred_x_start'])
@@ -713,3 +721,170 @@ def forward(self, img, *args, **kwargs):
 
         img = self.normalize(img)
         return self.p_losses(img, t, *args, **kwargs)
+
+# trainer class
+
+class Trainer1D(object):
+    def __init__(
+        self,
+        diffusion_model: GaussianDiffusion1D,
+        dataset: Dataset,
+        *,
+        train_batch_size = 16,
+        gradient_accumulate_every = 1,
+        train_lr = 1e-4,
+        train_num_steps = 100000,
+        ema_update_every = 10,
+        ema_decay = 0.995,
+        adam_betas = (0.9, 0.99),
+        save_and_sample_every = 1000,
+        num_samples = 25,
+        results_folder = './results',
+        amp = False,
+        fp16 = False,
+        split_batches = True,
+    ):
+        super().__init__()
+
+        # accelerator
+
+        self.accelerator = Accelerator(
+            split_batches = split_batches,
+            mixed_precision = 'fp16' if fp16 else 'no'
+        )
+
+        self.accelerator.native_amp = amp
+
+        # model
+
+        self.model = diffusion_model
+        self.channels = diffusion_model.channels
+
+        # sampling and training hyperparameters
+
+        assert has_int_squareroot(num_samples), 'number of samples must have an integer square root'
+        self.num_samples = num_samples
+        self.save_and_sample_every = save_and_sample_every
+
+        self.batch_size = train_batch_size
+        self.gradient_accumulate_every = gradient_accumulate_every
+
+        self.train_num_steps = train_num_steps
+
+        # dataset and dataloader
+
+        dl = DataLoader(dataset, batch_size = train_batch_size, shuffle = True, pin_memory = True, num_workers = cpu_count())
+
+        dl = self.accelerator.prepare(dl)
+        self.dl = cycle(dl)
+
+        # optimizer
+
+        self.opt = Adam(diffusion_model.parameters(), lr = train_lr, betas = adam_betas)
+
+        # for logging results in a folder periodically
+
+        if self.accelerator.is_main_process:
+            self.ema = EMA(diffusion_model, beta = ema_decay, update_every = ema_update_every)
+            self.ema.to(self.device)
+
+        self.results_folder = Path(results_folder)
+        self.results_folder.mkdir(exist_ok = True)
+
+        # step counter state
+
+        self.step = 0
+
+        # prepare model, dataloader, optimizer with accelerator
+
+        self.model, self.opt = self.accelerator.prepare(self.model, self.opt)
+
+    @property
+    def device(self):
+        return self.accelerator.device
+
+    def save(self, milestone):
+        if not self.accelerator.is_local_main_process:
+            return
+
+        data = {
+            'step': self.step,
+            'model': self.accelerator.get_state_dict(self.model),
+            'opt': self.opt.state_dict(),
+            'ema': self.ema.state_dict(),
+            'scaler': self.accelerator.scaler.state_dict() if exists(self.accelerator.scaler) else None,
+            'version': __version__
+        }
+
+        torch.save(data, str(self.results_folder / f'model-{milestone}.pt'))
+
+    def load(self, milestone):
+        accelerator = self.accelerator
+        device = accelerator.device
+
+        data = torch.load(str(self.results_folder / f'model-{milestone}.pt'), map_location=device)
+
+        model = self.accelerator.unwrap_model(self.model)
+        model.load_state_dict(data['model'])
+
+        self.step = data['step']
+        self.opt.load_state_dict(data['opt'])
+        if self.accelerator.is_main_process:
+            self.ema.load_state_dict(data["ema"])
+
+        if 'version' in data:
+            print(f"loading from version {data['version']}")
+
+        if exists(self.accelerator.scaler) and exists(data['scaler']):
+            self.accelerator.scaler.load_state_dict(data['scaler'])
+
+    def train(self):
+        accelerator = self.accelerator
+        device = accelerator.device
+
+        with tqdm(initial = self.step, total = self.train_num_steps, disable = not accelerator.is_main_process) as pbar:
+
+            while self.step < self.train_num_steps:
+
+                total_loss = 0.
+
+                for _ in range(self.gradient_accumulate_every):
+                    data = next(self.dl).to(device)
+
+                    with self.accelerator.autocast():
+                        loss = self.model(data)
+                        loss = loss / self.gradient_accumulate_every
+                        total_loss += loss.item()
+
+                    self.accelerator.backward(loss)
+
+                accelerator.clip_grad_norm_(self.model.parameters(), 1.0)
+                pbar.set_description(f'loss: {total_loss:.4f}')
+
+                accelerator.wait_for_everyone()
+
+                self.opt.step()
+                self.opt.zero_grad()
+
+                accelerator.wait_for_everyone()
+
+                self.step += 1
+                if accelerator.is_main_process:
+                    self.ema.update()
+
+                    if self.step != 0 and self.step % self.save_and_sample_every == 0:
+                        self.ema.ema_model.eval()
+
+                        with torch.no_grad():
+                            milestone = self.step // self.save_and_sample_every
+                            batches = num_to_groups(self.num_samples, self.batch_size)
+                            all_samples_list = list(map(lambda n: self.ema.ema_model.sample(batch_size=n), batches))
+                        #
+                        all_samples = torch.cat(all_samples_list, dim = 0)
+                        #
+                        torch.save(all_samples, str(self.results_folder / f'sample-{milestone}.png'))
+                        self.save(milestone)
+
+                pbar.update(1)
+
+        accelerator.print('training complete')