From cf2272c754a2efd5b1b84c59304a6cb7e0418b1c Mon Sep 17 00:00:00 2001
From: gdamms <damguillotin@gmail.com>
Date: Fri, 3 May 2024 16:15:34 +0200
Subject: paper was wrong?

---
 test.py | 405 +++++++++++++++-------------------------------------------------
 1 file changed, 95 insertions(+), 310 deletions(-)

(limited to 'test.py')

diff --git a/test.py b/test.py
index 79dd10b..d1db16d 100644
--- a/test.py
+++ b/test.py
@@ -1,310 +1,95 @@
-from typing import Tuple, Optional
-
-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
-
-from trainer import Trainer
-
-
-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)
-
-
-class UNet(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-        # 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)
-
-        # 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
-
-        # 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
-
-        # 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
-
-        # 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
-
-        # 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
-
-        # 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
-
-    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)
-
-        # Input (label)
-        y = self.inconv2(y)
-        y = y.view(-1, 1, 7, 7)
-
-        # Encoder
-        x = F.relu(self.e11(x))
-        x1 = F.relu(self.e12(x))
-        x = self.pool1(x1)
-
-        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')
+import numpy as np
+import matplotlib.pyplot as plt
+
+A = np.random.uniform(0.1, 0.9)
+B = np.random.uniform(0.1, 0.9)
+
+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 q_xt_xt_1(x, t):
+    alpha = ALPHA[t]
+    mean = np.sqrt(alpha) * x
+    std = 1 - alpha
+    return np.random.normal(mean, std)
+
+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)
+
+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)
+
+T = 100
+BETA = np.concatenate(([0], np.linspace(1e-4, 2e-2, T)))
+ALPHA = 1 - BETA
+ALPHA_BAR = np.cumprod(ALPHA)
+
+N = int(1e6)
+x0 = 1
+
+xs_implicit = np.array([x0] * N)
+for t in range(1, T+1):
+    xs_implicit = q_xt_xt_1(xs_implicit, t)
+
+xs_explicit_paper = q_xt_x0_paper(np.array([x0] * N), T)
+xs_explicit_damien = q_xt_x0_damien(np.array([x0] * N), T)
+
+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()
+
+xs_implicit = np.array([x0] * N)
+for t in range(1, T+1):
+    xs_implicit = q_xt_xt_1_simple(xs_implicit, t)
+
+xs_explicit_damien = q_xt_x0_simple_damien(np.array([x0] * N), T)
+
+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()
+
+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)
+
+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()
-- 
cgit v1.3.1