From 8c3609a6e3c216264e110c2019e61c83dafad9f5 Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Mon, 20 Jun 2022 13:17:51 -0700 Subject: [PATCH] move EMA logic out of the repository for clarity --- .../denoising_diffusion_pytorch.py | 47 +++++-------------- setup.py | 3 +- 2 files changed, 13 insertions(+), 37 deletions(-) diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 969c2bf..96222c5 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -19,6 +19,8 @@ from tqdm import tqdm from einops import rearrange, reduce from einops.layers.torch import Rearrange +from ema_pytorch import EMA + # helpers functions def exists(x): @@ -50,21 +52,6 @@ def unnormalize_to_zero_to_one(t): # small helper modules -class EMA(): - def __init__(self, beta): - super().__init__() - self.beta = beta - - def update_model_average(self, ma_model, current_model): - for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()): - old_weight, up_weight = ma_params.data, current_params.data - ma_params.data = self.update_average(old_weight, up_weight) - - def update_average(self, old, new): - if old is None: - return new - return old * self.beta + (1 - self.beta) * new - class Residual(nn.Module): def __init__(self, fn): super().__init__() @@ -612,8 +599,7 @@ class Trainer(object): self.image_size = diffusion_model.image_size self.model = diffusion_model - self.ema = EMA(ema_decay) - self.ema_model = copy.deepcopy(self.model) + self.ema = EMA(diffusion_model, beta = ema_decay) self.update_ema_every = update_ema_every self.step_start_ema = step_start_ema @@ -636,22 +622,11 @@ class Trainer(object): self.results_folder = Path(results_folder) self.results_folder.mkdir(exist_ok = True) - self.reset_parameters() - - def reset_parameters(self): - self.ema_model.load_state_dict(self.model.state_dict()) - - def step_ema(self): - if self.step < self.step_start_ema: - self.reset_parameters() - return - self.ema.update_model_average(self.ema_model, self.model) - def save(self, milestone): data = { 'step': self.step, 'model': self.model.state_dict(), - 'ema': self.ema_model.state_dict(), + 'ema': self.ema.state_dict(), 'scaler': self.scaler.state_dict() } torch.save(data, str(self.results_folder / f'model-{milestone}.pt')) @@ -661,7 +636,7 @@ class Trainer(object): self.step = data['step'] self.model.load_state_dict(data['model']) - self.ema_model.load_state_dict(data['ema']) + self.ema.load_state_dict(data['ema']) self.scaler.load_state_dict(data['scaler']) def train(self): @@ -681,15 +656,15 @@ class Trainer(object): self.scaler.update() self.opt.zero_grad() - if self.step % self.update_ema_every == 0: - self.step_ema() + self.ema.update() if self.step != 0 and self.step % self.save_and_sample_every == 0: - self.ema_model.eval() + self.ema.ema_model.eval() + with torch.no_grad(): + milestone = self.step // self.save_and_sample_every + batches = num_to_groups(36, self.batch_size) + all_images_list = list(map(lambda n: self.ema.ema_model.sample(batch_size=n), batches)) - milestone = self.step // self.save_and_sample_every - batches = num_to_groups(36, self.batch_size) - all_images_list = list(map(lambda n: self.ema_model.sample(batch_size=n), batches)) all_images = torch.cat(all_images_list, dim=0) utils.save_image(all_images, str(self.results_folder / f'sample-{milestone}.png'), nrow = 6) self.save(milestone) diff --git a/setup.py b/setup.py index d3d6ede..e021727 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.20.2', + version = '0.21.0', license='MIT', description = 'Denoising Diffusion Probabilistic Models - Pytorch', author = 'Phil Wang', @@ -16,6 +16,7 @@ setup( ], install_requires=[ 'einops', + 'ema-pytorch', 'pillow', 'torch', 'torchvision',