From c9855c8c7450c69c198b8f0642515c893d9f2644 Mon Sep 17 00:00:00 2001 From: Phil Wang Date: Sat, 5 Sep 2020 13:55:09 -0700 Subject: [PATCH] complete model portion --- denoising_diffusion_pytorch/__init__.py | 2 +- .../denoising_diffusion_pytorch.py | 178 +++++++++++++++++- setup.py | 5 +- 3 files changed, 181 insertions(+), 4 deletions(-) diff --git a/denoising_diffusion_pytorch/__init__.py b/denoising_diffusion_pytorch/__init__.py index 97f6d91..851bf72 100644 --- a/denoising_diffusion_pytorch/__init__.py +++ b/denoising_diffusion_pytorch/__init__.py @@ -1 +1 @@ -from denoising_diffusion_pytorch.denoising_diffusion_pytorch import DenoisingDiffusion +from denoising_diffusion_pytorch.denoising_diffusion_pytorch import DenoisingDiffusion, Unet diff --git a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py index 60558a2..b2849fd 100644 --- a/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py +++ b/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py @@ -1,11 +1,187 @@ +import math import torch from torch import nn, einsum import torch.nn.functional as F + +import numpy as np from einops import rearrange +# helper models + +class Residual(nn.Module): + def __init__(self, fn): + super().__init__() + self.fn = fn + + def forward(self, x, *args, **kwargs): + return self.fn(x, *args, **kwargs) + x + +class SinusoidalPosEmb(nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + + def forward(self, x): + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim) * -emb) + emb = x[:, None] * emb[None, :] + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + +class ResnetBlock(nn.Module): + def __init__(self, dim, out_dim, *, time_emb_dim, groups = 32): + super().__init__() + self.mlp = nn.Sequential( + Mish(), + nn.Linear(time_emb_dim, out_dim) + ) + + self.block1 = nn.Sequential( + nn.Conv2d(dim, out_dim, 3, padding=1), + nn.GroupNorm(groups, out_dim), + Mish() + ) + + self.block2 = nn.Sequential( + nn.Conv2d(out_dim, out_dim, 3, padding=1), + nn.GroupNorm(groups, out_dim), + Mish() + ) + + self.res_conv = nn.Conv2d(dim, out_dim, 1) if dim != out_dim else nn.Identity() + + def forward(self, x, time_emb): + h = self.block1(x) + h += self.mlp(time_emb)[:, :, None, None] + h = self.block2(h) + return h + self.res_conv(x) + +class Mish(nn.Module): + def forward(self, x): + return x * torch.tanh(F.softplus(x)) + +class Upsample(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.ConvTranspose2d(dim, dim, 4, 2, 1) + + def forward(self, x): + return self.conv(x) + +class Downsample(nn.Module): + def __init__(self, dim): + super().__init__() + self.conv = nn.Conv2d(dim, dim, 3, 2, 1) + + def forward(self, x): + return self.conv(x) + +class Rezero(nn.Module): + def __init__(self, dim): + super().__init__() + self.g = nn.Parameter(torch.zeros(1)) + + def forward(self, x): + return x * self.g + +class LinearAttention(nn.Module): + def __init__(self, dim, heads = 8, dim_head = 32): + super().__init__() + self.heads = heads + hidden_dim = dim_head * heads + self.to_qkv = nn.Conv2d(dim, hidden_dim, 1, bias = False) + self.to_out = nn.Conv2d(hidden_dim, dim, 1) + + def forward(self, x): + b, c, h, w = x.shape + qkv = self.to_qkv(x) + q, k, v = rearrange(qkv, 'b (qkv heads c) h w -> qkv b heads c (h w)', heads = self.heads) + q = q.softmax(dim=-2) + k = k.softmax(dim=-1) + context = torch.einsum('bhdn,bhen->bhde', k, v) + out = torch.einsum('bhde,bhdn->bhen', context, q) + out = rearrange(out, 'b heads c (h w) -> b (heads c) h w', heads=self.heads, h=h, w=w) + return self.to_out(out) + +# model + +class Unet(nn.Module): + def __init__(self, dim, dim_mults=(1, 2, 4, 8), groups = 32): + super().__init__() + dims = [3, *map(lambda m: dim * m, dim_mults)] + in_out = list(zip(dims[:-1], dims[1:])) + + self.time_pos_emb = SinusoidalPosEmb(dim) + self.mlp = nn.Sequential( + nn.Linear(dim, dim * 4), + Mish(), + nn.Linear(dim * 4, dim) + ) + + self.downs = nn.ModuleList([]) + self.ups = nn.ModuleList([]) + num_resolutions = len(in_out) + + for ind, (dim_in, dim_out) in enumerate(in_out): + is_last = ind >= (num_resolutions - 1) + + self.downs.append(nn.ModuleList([ + ResnetBlock(dim_in, dim_out, time_emb_dim = dim), + Residual(Rezero(LinearAttention(dim_out))), + Downsample(dim_out) if not is_last else nn.Identity() + ])) + + mid_dim = dims[-1] + self.mid_block1 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim) + self.mid_attn = Residual(Rezero(LinearAttention(mid_dim))) + self.mid_block2 = ResnetBlock(mid_dim, mid_dim, time_emb_dim = dim) + + for ind, (dim_in, dim_out) in enumerate(reversed(in_out[1:])): + is_last = ind >= (num_resolutions - 1) + + self.ups.append(nn.ModuleList([ + ResnetBlock(dim_out * 2, dim_in, time_emb_dim = dim), + Residual(Rezero(LinearAttention(dim_in))), + Upsample(dim_in) if not is_last else nn.Identity() + ])) + + self.final_conv = nn.Sequential( + nn.Conv2d(dim, dim, 3, padding = 1), + nn.GroupNorm(groups, dim), + Mish(), + nn.Conv2d(dim, 3, 1) + ) + + def forward(self, x, time): + t = self.time_pos_emb(time) + t = self.mlp(t) + + h = [] + + for resnet, attn, downsample in self.downs: + x = resnet(x, t) + x = attn(x) + h.append(x) + x = downsample(x) + + x = self.mid_block1(x, t) + x = self.mid_attn(x) + x = self.mid_block2(x, t) + + for resnet, attn, upsample in self.ups: + x = torch.cat((x, h.pop()), dim=1) + x = resnet(x, t) + x = attn(x) + x = upsample(x) + + return self.final_conv(x) + +# gaussian diffusion trainer class + class DenoisingDiffusion(nn.Module): def __init__(self): super().__init__() def forward(self, x): - return x \ No newline at end of file + return x diff --git a/setup.py b/setup.py index 2300d1d..3d86987 100644 --- a/setup.py +++ b/setup.py @@ -14,8 +14,9 @@ setup( 'generative models' ], install_requires=[ - 'torch', - 'einops' + 'einops', + 'numpy', + 'torch' ], classifiers=[ 'Development Status :: 4 - Beta',