diff options
Diffstat (limited to 'test.py')
| -rw-r--r-- | test.py | 371 |
1 files changed, 78 insertions, 293 deletions
@@ -1,310 +1,95 @@ -from typing import Tuple, Optional +import numpy as np +import matplotlib.pyplot as plt -import torch -import torch.nn.functional as F -import torch.utils.data -from torch import nn -from torch.utils.data import DataLoader, Dataset -from torchvision import datasets, transforms -from matplotlib import pyplot as plt -from rich.progress import track +A = np.random.uniform(0.1, 0.9) +B = np.random.uniform(0.1, 0.9) -from trainer import Trainer +def q_xt_xt_1_simple(x, t): + mean = A * x + std = B + return np.random.normal(mean, std) +def q_xt_x0_simple_damien(x, t): + mean = A ** t * x + std = np.sqrt(sum([A ** (2*i) for i in range(t)])) * B + return np.random.normal(mean, std) -def gather(consts: torch.Tensor, t: torch.Tensor): - """Gather consts for $t$ and reshape to feature map shape""" - c = consts.gather(-1, t) - return c.reshape(-1, 1, 1, 1) +def q_xt_xt_1(x, t): + alpha = ALPHA[t] + mean = np.sqrt(alpha) * x + std = 1 - alpha + return np.random.normal(mean, std) -class UNet(nn.Module): - def __init__(self): - super().__init__() +def q_xt_x0_paper(x, t): + alpha_bar = ALPHA_BAR[t] + mean = np.sqrt(alpha_bar) * x + std = 1 - alpha_bar + return np.random.normal(mean, std) - # Input - # The input to the model is a 11 vector which represents the desired label with the contextual information. - # The Input is passed through layers to generate two feature maps of size 7x7. - # ------- - # input: 1 (diffu step) and 10 (label) - self.inconv1 = nn.Linear(1, 7 * 7) - self.inconv2 = nn.Linear(10, 7 * 7) +def q_xt_x0_damien(x, t): + alpha_bar = ALPHA_BAR[t] + cum_sq_sum = sum([np.prod(ALPHA[s+2:t+1]) * (1 - ALPHA[s+1])**2 for s in range(t)]) + mean = np.sqrt(alpha_bar) * x + std = np.sqrt(cum_sq_sum) + return np.random.normal(mean, std) - # Encoder - # In the encoder, convolutional layers with the Conv2d function are used to extract features from the input image. - # Each block in the encoder consists of two convolutional layers followed by a max-pooling layer, - # with the exception of the last block which does not include a max-pooling layer. - # ------- - # input: 28x28x1 - self.e11 = nn.Conv2d(1, 64, kernel_size=3, - padding=1) # output: 28x28x64 - self.e12 = nn.Conv2d(64, 64, kernel_size=3, - padding=1) # output: 28x28x64 - self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) # output: 14x14x64 +T = 100 +BETA = np.concatenate(([0], np.linspace(1e-4, 2e-2, T))) +ALPHA = 1 - BETA +ALPHA_BAR = np.cumprod(ALPHA) - # input: 14x14x64 - self.e21 = nn.Conv2d(64, 128, kernel_size=3, - padding=1) # output: 14x14x128 - self.e22 = nn.Conv2d(128, 128, kernel_size=3, - padding=1) # output: 14x14x128 - self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) # output: 7x7x128 +N = int(1e6) +x0 = 1 - # input: 7x7x128 - self.e31 = nn.Conv2d(130, 256, kernel_size=3, - padding=1) # output: 7x7x256 - self.e32 = nn.Conv2d(258, 256, kernel_size=3, - padding=1) # output: 7x7x256 +xs_implicit = np.array([x0] * N) +for t in range(1, T+1): + xs_implicit = q_xt_xt_1(xs_implicit, t) - # Decoder - # In the decoder, the output of the encoder is upsampled using the ConvTranspose2d function. - # Each block in the decoder consists of two convolutional layers followed by an upsampling layer, - # with the exception of the last block which does not include an upsampling layer. - # ------- - # input: 7x7x256 - self.upconv1 = nn.ConvTranspose2d( - 256, 128, kernel_size=2, stride=2) # output: 14x14x128 - self.d11 = nn.Conv2d(256, 128, kernel_size=3, - padding=1) # output: 14x14x(128x2) - self.d12 = nn.Conv2d(128, 128, kernel_size=3, - padding=1) # output: 14x14x128 +xs_explicit_paper = q_xt_x0_paper(np.array([x0] * N), T) +xs_explicit_damien = q_xt_x0_damien(np.array([x0] * N), T) - # input: 14x14x128 - self.upconv2 = nn.ConvTranspose2d( - 128, 64, kernel_size=2, stride=2) # output: 28x28x64 - self.d21 = nn.Conv2d(128, 64, kernel_size=3, - padding=1) # output: 28x28x(64x2) - self.d22 = nn.Conv2d(64, 64, kernel_size=3, - padding=1) # output: 28x28x64 +plt.figure() +bins = np.linspace(min( + xs_implicit.min(), + xs_explicit_paper.min(), + xs_explicit_damien.min(), + ), max( + xs_implicit.max(), + xs_explicit_paper.max(), + xs_explicit_damien.max(), + ), 100) +plt.hist(xs_implicit, bins=bins, alpha=0.5, label="q_xt_xt_1") +plt.hist(xs_explicit_paper, bins=bins, alpha=0.5, label="q_xt_x0_paper") +plt.hist(xs_explicit_damien, bins=bins, alpha=0.5, label="q_xt_x0_damien") +plt.legend() +plt.savefig("q_xt_xt_1_vs_q_xt_x0.tmp.png") +plt.show() - # Output - # The output of the decoder is passed through a convolutional layer with the Conv2d function to obtain the final output. - # ------- - # input: 28x28x64 - self.outconv = nn.Conv2d(64, 1, kernel_size=1) # output: 28x28x2 +xs_implicit = np.array([x0] * N) +for t in range(1, T+1): + xs_implicit = q_xt_xt_1_simple(xs_implicit, t) - def forward(self, x, t, y): - # Input (diffusion step) - t = t.unsqueeze_(-1) - t = t.to(torch.float32) - t = self.inconv1(t) - t = t.view(-1, 1, 7, 7) +xs_explicit_damien = q_xt_x0_simple_damien(np.array([x0] * N), T) - # Input (label) - y = self.inconv2(y) - y = y.view(-1, 1, 7, 7) +plt.figure() +bins = np.linspace(min(xs_implicit.min(), xs_explicit_damien.min()), max(xs_implicit.max(), xs_explicit_damien.max()), 100) +plt.hist(xs_implicit, bins=bins, alpha=0.5, label="q_xt_xt_1_simple") +plt.hist(xs_explicit_damien, bins=bins, alpha=0.5, label="q_xt_x0_simple_damien") +plt.legend() +plt.savefig("q_xt_xt_1_simple_vs_q_xt_x0_simple.tmp.png") +plt.show() - # Encoder - x = F.relu(self.e11(x)) - x1 = F.relu(self.e12(x)) - x = self.pool1(x1) +A1, B1, A2, B2, A3, B3 = np.random.uniform(0, 1, 6) +x1 = np.random.normal(A1, B1, N) +x2 = np.random.normal(A2 * x1, B2, N) +x3 = np.random.normal(A3 * x2, B3, N) +x3_ = np.random.normal(A1 * A2 * A3, np.sqrt(A3**2 * A2**2 * B1**2 + A3**2 * B2**2 + B3**2), N) - x = F.relu(self.e21(x)) - x2 = F.relu(self.e22(x)) - x = self.pool2(x2) - - x = torch.cat([x, t, y], dim=1) - x = F.relu(self.e31(x)) - x = torch.cat([x, t, y], dim=1) - x = F.relu(self.e32(x)) - - # Decoder - x = self.upconv1(x) - x = torch.cat([x, x], dim=1) - x = F.relu(self.d11(x)) - x = F.relu(self.d12(x)) - - x = self.upconv2(x) - x = torch.cat([x, x1], dim=1) - x = F.relu(self.d21(x)) - x = F.relu(self.d22(x)) - - # Output - x = self.outconv(x) - - return x - - -class MNISTDiffusionDataset(Dataset): - def __init__(self, train=True): - super().__init__() - self.mnist_data = datasets.MNIST( - root='./data', - train=train, - download=True, - transform=transforms.ToTensor(), - ) - - def __len__(self): - return len(self.mnist_data) - - def __getitem__(self, index): - # Get the image and the label. - img, label = self.mnist_data[index] - prompt = torch.nn.functional.one_hot( - torch.tensor(label), 10).to(torch.float32) - return img.to(device), prompt.to(device) - - -class DenoiseDiffusion: - """ - ## Denoise Diffusion - """ - - def __init__(self, eps_model: nn.Module, n_steps: int, device: torch.device): - """ - * `eps_model` is $\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$ model - * `n_steps` is $t$ - * `device` is the device to place constants on - """ - super().__init__() - self.eps_model = eps_model - - # Create $\beta_1, \dots, \beta_T$ linearly increasing variance schedule - self.beta = torch.linspace(0.0001, 0.02, n_steps).to(device) - - # $\alpha_t = 1 - \beta_t$ - self.alpha = 1. - self.beta - # $\bar\alpha_t = \prod_{s=1}^t \alpha_s$ - self.alpha_bar = torch.cumprod(self.alpha, dim=0) - # $T$ - self.n_steps = n_steps - # $\sigma^2 = \beta$ - self.sigma2 = self.beta - - def q_xt_x0(self, x0: torch.Tensor, t: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """ - #### Get $q(x_t|x_0)$ distribution - - \begin{align} - q(x_t|x_0) &= \mathcal{N} \Big(x_t; \sqrt{\bar\alpha_t} x_0, (1-\bar\alpha_t) \mathbf{I} \Big) - \end{align} - """ - - # [gather](utils.html) $\alpha_t$ and compute $\sqrt{\bar\alpha_t} x_0$ - mean = gather(self.alpha_bar, t) ** 0.5 * x0 - # $(1-\bar\alpha_t) \mathbf{I}$ - var = 1 - gather(self.alpha_bar, t) - # - return mean, var - - def q_sample(self, x0: torch.Tensor, t: torch.Tensor, eps: Optional[torch.Tensor] = None): - """ - #### Sample from $q(x_t|x_0)$ - - \begin{align} - q(x_t|x_0) &= \mathcal{N} \Big(x_t; \sqrt{\bar\alpha_t} x_0, (1-\bar\alpha_t) \mathbf{I} \Big) - \end{align} - """ - - # $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ - if eps is None: - eps = torch.randn_like(x0) - - # get $q(x_t|x_0)$ - mean, var = self.q_xt_x0(x0, t) - # Sample from $q(x_t|x_0)$ - return mean + (var ** 0.5) * eps - - def p_sample(self, xt: torch.Tensor, t: torch.Tensor, prompt: Optional[torch.Tensor] = None): - """ - #### Sample from $\textcolor{lightgreen}{p_\theta}(x_{t-1}|x_t)$ - - \begin{align} - \textcolor{lightgreen}{p_\theta}(x_{t-1} | x_t) &= \mathcal{N}\big(x_{t-1}; - \textcolor{lightgreen}{\mu_\theta}(x_t, t), \sigma_t^2 \mathbf{I} \big) \\ - \textcolor{lightgreen}{\mu_\theta}(x_t, t) - &= \frac{1}{\sqrt{\alpha_t}} \Big(x_t - - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big) - \end{align} - """ - - # $\textcolor{lightgreen}{\epsilon_\theta}(x_t, t)$ - eps_theta = self.eps_model(xt, t, prompt) - # [gather](utils.html) $\bar\alpha_t$ - alpha_bar = gather(self.alpha_bar, t) - # $\alpha_t$ - alpha = gather(self.alpha, t) - # $\frac{\beta}{\sqrt{1-\bar\alpha_t}}$ - eps_coef = (1 - alpha) / (1 - alpha_bar) ** .5 - # $$\frac{1}{\sqrt{\alpha_t}} \Big(x_t - - # \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\textcolor{lightgreen}{\epsilon_\theta}(x_t, t) \Big)$$ - mean = 1 / (alpha ** 0.5) * (xt - eps_coef * eps_theta) - # $\sigma^2$ - var = gather(self.sigma2, t) - - # $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ - eps = torch.randn(xt.shape, device=xt.device) - # Sample - return mean + (var ** .5) * eps - - def loss(self, x0: torch.Tensor, prompt: Optional[torch.Tensor] = None, noise: Optional[torch.Tensor] = None): - """ - #### Simplified Loss - - $$L_{\text{simple}}(\theta) = \mathbb{E}_{t,x_0, \epsilon} \Bigg[ \bigg\Vert - \epsilon - \textcolor{lightgreen}{\epsilon_\theta}(\sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon, t) - \bigg\Vert^2 \Bigg]$$ - """ - # Get batch size - batch_size = x0.shape[0] - # Get random $t$ for each sample in the batch - t = torch.randint(0, self.n_steps, (batch_size,), - device=x0.device, dtype=torch.long) - - # $\epsilon \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ - if noise is None: - noise = torch.randn_like(x0) - - # Sample $x_t$ for $q(x_t|x_0)$ - xt = self.q_sample(x0, t, eps=noise) - # Get $\textcolor{lightgreen}{\epsilon_\theta}(\sqrt{\bar\alpha_t} x_0 + \sqrt{1-\bar\alpha_t}\epsilon, t)$ - if prompt is None: - eps_theta = self.eps_model(xt, t) - else: - eps_theta = self.eps_model(xt, t, prompt) - - # MSE loss - return F.mse_loss(noise, eps_theta) - - -diffu_steps = 64 - -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') -model = UNet().to(device) -model = torch.load('model.pt').to(device) -ddpm = DenoiseDiffusion(model, diffu_steps, device) - -ds = MNISTDiffusionDataset() -dl = DataLoader(ds, batch_size=128, shuffle=True) -opti = torch.optim.Adam(model.parameters(), lr=1e-5) - - -for epoch in range(3): - for x0, prompt in track(dl): - loss = ddpm.loss(x0, prompt=prompt) - opti.zero_grad() - loss.backward() - opti.step() - - print(f'Epoch {epoch}: {loss.item()}') - -torch.save(model, 'model.pt') - -n = 1 -fig = plt.figure(figsize=(10, n)) -x = torch.randn(10 * n, 1, 28, 28, device=device) -prompt = torch.nn.functional.one_hot( - torch.tensor([range(10)] * n), num_classes=10).to(device, torch.float32) -prompt = prompt.reshape(-1, 10) -for i in track(range(diffu_steps)): - for j in range(10 * n): - t = diffu_steps - i - 1 - t = torch.tensor(t, device=device) - x[j] = ddpm.p_sample(x[j:j+1], t, prompt[j]) - -for i in range(10 * n): - plt.subplot(n, 10, i + 1) - plt.axis('off') - plt.imshow(x[i, 0].cpu().detach().numpy()) -plt.tight_layout() -plt.savefig('test.png') +plt.figure() +bins = np.linspace(min(x3.min(), x3_.min()), max(x3.max(), x3_.max()), 100) +plt.hist(x3, bins=bins, alpha=0.5, label="normal") +plt.hist(x3_, bins=bins, alpha=0.5, label="product") +plt.legend() +plt.savefig("product_normal.tmp.png") +plt.show() |