diff options
Diffstat (limited to 'test.py')
| -rw-r--r-- | test.py | 310 |
1 files changed, 310 insertions, 0 deletions
@@ -0,0 +1,310 @@ +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') |