diff options
| -rw-r--r-- | main.py | 238 | ||||
| -rw-r--r-- | model.pt | bin | 0 -> 7499403 bytes | |||
| -rw-r--r-- | test.png | bin | 0 -> 58259 bytes | |||
| -rw-r--r-- | test.py | 310 | ||||
| -rw-r--r-- | trainer.py | 303 |
5 files changed, 751 insertions, 100 deletions
@@ -17,13 +17,13 @@ class UNet(nn.Module): super().__init__() # Input - # The input to the model is a 10 vector which represents the input image. + # 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 a 1x28x28, 1x14x14, 1x7x7 tensor. # ------- - # input: 1x10 - self.inconv1 = nn.Linear(10, 28 * 28) - self.inconv2 = nn.Linear(10, 14 * 14) - self.inconv3 = nn.Linear(10, 7 * 7) + # input: 1x11 + self.inconv1 = nn.Linear(11, 28 * 28) + self.inconv2 = nn.Linear(11, 14 * 14) + self.inconv3 = nn.Linear(11, 7 * 7) # Encoder # In the encoder, convolutional layers with the Conv2d function are used to extract features from the input image. @@ -31,18 +31,24 @@ class UNet(nn.Module): # 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.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.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(129, 256, kernel_size=3, padding=1) # output: 7x7x256 - self.e32 = nn.Conv2d(257, 256, kernel_size=3, padding=1) # output: 7x7x256 + self.e31 = nn.Conv2d(129, 256, kernel_size=3, + padding=1) # output: 7x7x256 + self.e32 = nn.Conv2d(257, 256, kernel_size=3, + padding=1) # output: 7x7x256 # Decoder # In the decoder, the output of the encoder is upsampled using the ConvTranspose2d function. @@ -50,20 +56,26 @@ class UNet(nn.Module): # 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 + 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 + 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: 28x28x1 + self.outconv = nn.Conv2d(64, 1, kernel_size=1) # output: 28x28x2 def forward(self, x, y): # Input @@ -101,48 +113,125 @@ class UNet(nn.Module): return x +DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + DIFFU_STEPS = 10 +BETA = torch.linspace(0.0001, 0.2, DIFFU_STEPS+1, device=DEVICE) +ALPHA = 1 - BETA +ALPHA_BAR = torch.cumprod(ALPHA, dim=0) +SIGMA2 = BETA + + +def q_xt_x0(x0, t): + alpha_bar = ALPHA_BAR[t] + mean = x0 * torch.sqrt(alpha_bar) + std = 1 - alpha_bar + return torch.distributions.Normal(mean, std) + + +def q_xt_xt_1(xt_1, t): + beta = BETA[t] + mean = torch.sqrt(1 - beta) * xt_1 + std = beta + return torch.distributions.Normal(mean, std) + + +def p_xt_1_xt(model, xt, vec): + t = vec[..., -1].to(dtype=torch.long) + eps_theta = model(xt, vec) + alpha_bar = ALPHA_BAR[t].unsqueeze( + -1).unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 28, 28) + alpha = ALPHA[t].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 28, 28) + beta = BETA[t].unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 28, 28) + eps_coef = (1 - alpha) / torch.sqrt(1 - alpha_bar) + # mean = 1 / torch.sqrt(alpha) * (xt - eps_coef * eps_theta) + mean = (xt - beta * eps_theta) + # std = torch.sqrt( + # SIGMA2[t]).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).repeat(1, 1, 28, 28) + std = beta + return torch.distributions.Normal(mean, std) class MNISTDiffusionDataset(Dataset): def __init__(self, train=True): super().__init__() - self.mnist_data = datasets.MNIST(root='./data', train=train, download=True, transform=transforms.ToTensor()) + self.mnist_data = datasets.MNIST( + root='./data', + train=train, + download=True, + transform=transforms.ToTensor(), + ) def __getitem__(self, index): # Get the image and the label. img, label = self.mnist_data[index] + img = img.to(DEVICE) # Add noise to the image. - noise = np.random.normal(0, 1, (28, 28)) - alpha = np.random.uniform(1 / DIFFU_STEPS, 1.0) - - # The target is the image with the noise. - target = img * alpha + noise * (1 - alpha) - - # The input is the image with more noise. - input = img * (alpha - 1 / DIFFU_STEPS) + noise * (1 - alpha + 1 / DIFFU_STEPS) + t_1 = torch.randint(0, DIFFU_STEPS, (1,), device=DEVICE) + t = t_1 + 1 + xt_1 = q_xt_x0(img, t_1).sample() + eps = torch.distributions.Normal(0, 1).sample(img.shape).to(DEVICE) + xt = xt_1 * torch.sqrt(ALPHA[t]) + BETA[t] * eps # Convert the label to a one-hot vector. - vector = torch.nn.functional.one_hot(torch.tensor(label), num_classes=10) + vector = torch.nn.functional.one_hot( + torch.tensor(label), + num_classes=10, + ) - return (input.clone().detach().to(dtype=torch.float32), - vector.clone().detach().to(dtype=torch.float32), - target.clone().detach().to(dtype=torch.float32)) + # Add contextual information (t) to the label. + vector = torch.cat([vector, torch.tensor([t_1+1])]) + + return ( + xt.clone().detach().to(dtype=torch.float32, device=DEVICE), + vector.clone().detach().to(dtype=torch.float32, device=DEVICE), + ( + img.clone().detach().to(dtype=torch.float32, device=DEVICE), + eps, + t, + ), + ) def __len__(self): return len(self.mnist_data) +def loss_fn(y_pred, y_true): + x0, eps, t = y_true + return nn.MSELoss()(y_pred, eps) + + +mnist_data = datasets.MNIST( + root='./data', + train=True, + download=True, + transform=transforms.ToTensor(), +) +img, label = mnist_data[0] +img = img.to(DEVICE) +fig = plt.figure(figsize=(DIFFU_STEPS, 2)) +for t_1 in range(0, DIFFU_STEPS): + xt_1 = q_xt_x0(img, t_1).sample() + x_t = q_xt_xt_1(xt_1, t_1+1).sample() + ax = fig.add_subplot(2, DIFFU_STEPS, t_1 + 1) + ax.imshow(xt_1[0].cpu(), cmap='gray') + ax.axis('off') + ax = fig.add_subplot(2, DIFFU_STEPS, DIFFU_STEPS + t_1 + 1) + ax.imshow(x_t[0].cpu(), cmap='gray') + ax.axis('off') +fig.tight_layout() +fig.savefig('img.tmp.png') + + ############ # Training # ############ -# Define the device. -device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') +torch.multiprocessing.set_start_method('spawn') # Load the model. -model = UNet().to(device) +model = UNet().to(DEVICE) # model.load_state_dict(torch.load('model.pth')) # Define the optimizer. @@ -150,11 +239,13 @@ optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) # Define the training dataset. train_dataset = MNISTDiffusionDataset(train=True) -train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) +train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) trainer = Trainer() +criterion = loss_fn +epochs = 1 # Train the model. -trainer.train(model, train_loader, 2, optimizer, F.mse_loss) +trainer.train(model, train_loader, epochs, optimizer, criterion) # Save the model. torch.save(model.state_dict(), 'model.pth') @@ -164,8 +255,31 @@ torch.save(model.state_dict(), 'model.pth') # Evaluation # ############## +# for data in train_loader: +# input, vector, (x0, eps, t) = data +# eps_theta = model(input, vector) + +# fig = plt.figure(figsize=(2, 2)) +# ax = fig.add_subplot(2, 2, 1) +# ax.imshow(eps[0].cpu().transpose(0, 2).transpose(0, 1), cmap='gray') +# ax.axis('off') +# ax = fig.add_subplot(2, 2, 2) +# ax.imshow(eps_theta[0].cpu().detach().transpose( +# 0, 2).transpose(0, 1), cmap='gray') +# ax.axis('off') +# ax = fig.add_subplot(2, 2, 3) +# ax.imshow((eps[0] - eps_theta[0]).cpu().detach().transpose( +# 0, 2).transpose(0, 1), cmap='gray') +# ax.axis('off') +# ax = fig.add_subplot(2, 2, 4) +# ax.imshow(x0[0].cpu().transpose(0, 2).transpose(0, 1), cmap='gray') +# ax.axis('off') +# fig.tight_layout() +# fig.savefig('eps.tmp.png') +# exit() + # Load the model. -model = UNet().to(device) +model = UNet().to(DEVICE) model.load_state_dict(torch.load('model.pth')) # Set the model to evaluation mode. @@ -176,29 +290,35 @@ fig = plt.figure(figsize=(2 * 2 * n, 3 * 2)) gs = plt.GridSpec(nrows=3, ncols=2*2*n) for i in range(n): # Get the i-th input and its label. - input, vector, label = train_dataset[i] + input, vector, (x0, eps, t) = train_dataset[i] + + xt_1 = q_xt_x0(input, t-1).sample() + xt = q_xt_xt_1(xt_1, t).sample() + xt_1_pred = p_xt_1_xt( + model, + xt.unsqueeze(0), + vector.unsqueeze(0), + ).sample() # Plot the input. ax = fig.add_subplot(gs[0:1, 1 + 4 * i:3 + 4 * i]) - ax.imshow(input[0], cmap='gray') + ax.imshow(xt[0].cpu(), cmap='gray') ax.axis('off') # Plot the label. ax = fig.add_subplot(gs[1:2, 4 * i:2 + 4 * i]) - ax.imshow(label[0], cmap='gray') + ax.imshow(xt_1[0].cpu(), cmap='gray') ax.axis('off') - # Get the model output. - output = model(input.unsqueeze(0).to(device), vector.unsqueeze(0).to(device)) - # Plot the model output. ax = fig.add_subplot(gs[1:2, 2 + 4 * i:4 + 4 * i]) - ax.imshow(output[0, 0].cpu().detach(), cmap='gray') + ax.imshow(xt_1_pred[0, 0].cpu().detach(), cmap='gray') ax.axis('off') # Plot the difference between the label and the model output. ax = fig.add_subplot(gs[2:3, 1 + 4 * i:3 + 4 * i]) - ax.imshow((label - output[0, 0].cpu().detach())[0], cmap='coolwarm') + ax.imshow( + (xt_1.cpu() - xt_1_pred[0, 0].cpu().detach())[0], cmap='coolwarm') ax.axis('off') fig.tight_layout() @@ -208,13 +328,25 @@ fig.savefig('diff.tmp.png') # Plot the evolution of the noise. fig = plt.figure(figsize=(n, DIFFU_STEPS)) noises = np.random.normal(0, 1, (n, 1, 28, 28)) -noises = torch.Tensor(noises).to(device) -vector = torch.nn.functional.one_hot(torch.tensor(range(n)), num_classes=10).to(device) +noises = torch.Tensor(noises).to(DEVICE) +vector = torch.nn.functional.one_hot( + torch.tensor(range(n)), num_classes=10).to(DEVICE) vector = vector.clone().detach().to(dtype=torch.float32) # Apply the model multiple times. for i in range(DIFFU_STEPS): - noises = model(noises, vector) + t = DIFFU_STEPS - i - 1 + noises = p_xt_1_xt( + model, + noises, + torch.cat([ + vector, + torch.tensor([t] * n) + .unsqueeze(-1) + .to(device=DEVICE) + .to(dtype=torch.long), + ], dim=-1), + ).sample() for j in range(n): ax = fig.add_subplot(DIFFU_STEPS, n, i * n + j + 1) @@ -227,13 +359,17 @@ fig.savefig('diffu.tmp.png') # Plot bench of generated images. fig = plt.figure(figsize=(n, n)) noises = np.random.normal(0, 1, (n * n, 1, 28, 28)) -noises = torch.Tensor(noises).to(device) -vector = torch.nn.functional.one_hot(torch.tensor([range(n)] * n), num_classes=10).to(device) +noises = torch.Tensor(noises).to(DEVICE) +vector = torch.nn.functional.one_hot( + torch.tensor([range(n)] * n), num_classes=10).to(DEVICE) vector = vector.clone().detach().to(dtype=torch.float32) # Apply the model multiple times. for i in range(DIFFU_STEPS): - noises = model(noises, vector) + t = DIFFU_STEPS - i - 1 + noises = model(noises, torch.cat( + [vector, torch.tensor([[t] * n] * n).unsqueeze(-1) + .to(DEVICE)], dim=-1)) for i in range(n * n): ax = fig.add_subplot(n, n, i + 1) diff --git a/model.pt b/model.pt Binary files differnew file mode 100644 index 0000000..810cf1a --- /dev/null +++ b/model.pt diff --git a/test.png b/test.png Binary files differnew file mode 100644 index 0000000..b326a8e --- /dev/null +++ b/test.png @@ -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') @@ -12,7 +12,9 @@ class TrainProgress(rich.progress.Progress): def __init__( self: 'TrainProgress', nb_epochs: int, - epoch_size: int, + train_size: int, + val_size: int = 0, + test_size: int = 0, *columns: str | rich.progress.ProgressColumn, console: rich.progress.Console | None = None, auto_refresh: bool = True, @@ -29,7 +31,9 @@ class TrainProgress(rich.progress.Progress): Args: nb_epochs (int): The number of epochs. - epoch_size (int): The size of each epoch. + train_size (int): The size of each tain epoch. + val_size (int, optional): The size of each validation epoch. Defaults to 0. + test_size (int, optional): The size of the test epoch. Defaults to 0. *columns (str | rich.progress.ProgressColumn): The columns to display. console (rich.progress.Console, optional): The console to use. Defaults to None. auto_refresh (bool, optional): Whether to automatically refresh the progress bar. Defaults to True. @@ -43,7 +47,9 @@ class TrainProgress(rich.progress.Progress): expand (bool, optional): Whether to expand the progress bar. Defaults to False. """ self.nb_epochs = nb_epochs - self.epoch_size = epoch_size + self.train_size = train_size + self.val_size = val_size + self.test_size = test_size super().__init__( *columns, console=console, @@ -57,9 +63,17 @@ class TrainProgress(rich.progress.Progress): disable=disable, expand=expand, ) - self.epoch_tasks = [] - self.total_task = self.add_task("total", progress_type="total", total=nb_epochs*epoch_size) - self.values = {} + self.train_tasks = [] + self.val_tasks = [] + self.test_task = None + self.total_task = self.add_task( + "total", + progress_type="total", + total=nb_epochs * (train_size + val_size) + test_size, + ) + self.train_values = [] + self.val_values = [] + self.test_values = {} def get_renderables(self: 'TrainProgress'): """Override the default renderables to display the epoch number.""" @@ -68,58 +82,205 @@ class TrainProgress(rich.progress.Progress): # The total task. if task.fields.get("progress_type") == "total": self.columns = ( - f"Training:", + f"Working:", rich.progress.BarColumn(), - f"{len(self.epoch_tasks):{pad}}/{self.nb_epochs}", + f"{len(self.train_tasks):{pad}}/{self.nb_epochs}", "•", rich.progress.TimeRemainingColumn(), ) - # The epoch tasks. - if task.fields.get("progress_type") == "epoch": + # The train tasks. + if task.fields.get("progress_type") == "train": epoch_id = task.fields.get("epoch_id") self.columns = ( - f"Epoch {epoch_id:{pad}}:", + f"Train {epoch_id:{pad}}:", rich.progress.BarColumn(), f"{task.completed}/{task.total}", "•", rich.progress.TimeElapsedColumn(), '•', - ' | '.join(f"{key}: {value[-1]:.4f} " for key, value in self.values.items()), + ' | '.join( + f"{key}: {value[-1]:.4f}" for key, value in self.train_values[epoch_id-1].items()), + ) + + # The val tasks. + if task.fields.get("progress_type") == "val": + epoch_id = task.fields.get("epoch_id") + self.columns = ( + f"Val {epoch_id:{pad}}:", + rich.progress.BarColumn(), + f"{task.completed}/{task.total}", + "•", + rich.progress.TimeElapsedColumn(), + '•', + ' | '.join( + f"{key}: {value[-1]:.4f}" for key, value in self.val_values[epoch_id-1].items()), + ) + + # The test task. + if task.fields.get("progress_type") == "test": + self.columns = ( + f"Test:", + rich.progress.BarColumn(), + f"{task.completed}/{task.total}", + "•", + rich.progress.TimeElapsedColumn(), + '•', + ' | '.join( + f"{key}: {value[-1]:.4f}" for key, value in self.test_values.items()), ) yield self.make_tasks_table([task]) - def new_epoch(self: 'TrainProgress'): - """Create a new epoch task.""" - epoch_task = self.add_task( - "epoch", - progress_type="epoch", - epoch_id=len(self.epoch_tasks) + 1, - total=self.epoch_size, - ) - self.epoch_tasks.append(epoch_task) + def step_test(self: 'TrainProgress', count: int) -> bool: + """Advance the progress bar by the given number of steps. - def step(self: 'TrainProgress', count: int = 1): + Args: + count (int): The number of steps to advance the progress bar by. + + Returns: + bool: Whether step was successful. + """ + if len(self.train_tasks) < self.nb_epochs: + return False + + if self.tasks[self.train_tasks[-1]].completed < self.train_size: + return False + + if self.val_size > 0: + if len(self.val_tasks) < self.nb_epochs: + return False + + if self.tasks[self.val_tasks[-1]].completed < self.val_size: + return False + + if self.test_size == 0: + return False + + if self.test_task is None: + self.test_task = self.add_task( + "Test", + progress_type="test", + total=self.test_size, + ) + self.update(self.test_task, advance=count) + self.update(self.total_task, advance=count) + return True + + if self.test_task is not None: + self.update(self.test_task, advance=count) + self.update(self.total_task, advance=count) + return True + + def step_val(self: 'TrainProgress', count: int) -> bool: """Advance the progress bar by the given number of steps. Args: count (int): The number of steps to advance the progress bar by. + + Returns: + bool: Whether step was successful. """ - if len(self.epoch_tasks) == 0 or self.tasks[self.epoch_tasks[-1]].completed == self.epoch_size: - self.new_epoch() - self.update(self.epoch_tasks[-1], advance=count) + if len(self.train_tasks) == 0: + return False + + if self.tasks[self.train_tasks[-1]].completed < self.train_size: + return False + + if self.val_size == 0: + return False + + if len(self.val_tasks) == 0 or ( + len(self.val_tasks) < self.nb_epochs + and len(self.val_tasks) < len(self.train_tasks) + ): + self.val_values.append({}) + self.val_tasks.append(self.add_task( + f"Val {len(self.val_tasks)+1}", + progress_type="val", + epoch_id=len(self.val_tasks)+1, + total=self.val_size, + )) + self.update(self.val_tasks[-1], advance=count) + self.update(self.total_task, advance=count) + return True + + if self.tasks[self.val_tasks[-1]].completed < self.val_size: + self.update(self.val_tasks[-1], advance=count) + self.update(self.total_task, advance=count) + return True + + def step_train(self: 'TrainProgress', count: int) -> bool: + """Advance the progress bar by the given number of steps. + + Args: + count (int): The number of steps to advance the progress bar by. + + Returns: + bool: Whether step was successful. + """ + if len(self.train_tasks) == 0 or self.tasks[self.train_tasks[-1]].completed == self.train_size: + self.train_values.append({}) + self.train_tasks.append(self.add_task( + f"Train {len(self.train_tasks)+1}", + progress_type="train", + epoch_id=len(self.train_tasks)+1, + total=self.train_size, + )) + self.update(self.train_tasks[-1], advance=count) + self.update(self.total_task, advance=count) + return True + + self.update(self.train_tasks[-1], advance=count) self.update(self.total_task, advance=count) + return True + + def step(self: 'TrainProgress', count: int = 1): + """Advance the progress bar by the given number of steps. + + Args: + count (int): The number of steps to advance the progress bar by. + """ + if self.step_test(count): + return + + if self.step_val(count): + return + + if self.step_train(count): + return + + raise RuntimeError("Progress bar already finished.") + + def new_train_values(self: 'TrainProgress', values: dict[str, Any]): + """Update the progress bar with new values. + + Args: + values (dict[str, Any]): The new values. + """ + for key, value in values.items(): + current_value = self.train_values[-1].get(key, []) + self.train_values[-1][key] = current_value + [value] - def new_values(self: 'TrainProgress', **values: Any): + def new_val_values(self: 'TrainProgress', values: dict[str, Any]): """Update the progress bar with new values. Args: - **values (Any): The values to update the progress bar with. + values (dict[str, Any]): The new values. """ for key, value in values.items(): - current_value = self.values.get(key, []) - self.values[key] = current_value + [value] + current_value = self.val_values[-1].get(key, []) + self.val_values[-1][key] = current_value + [value] + + def new_test_values(self: 'TrainProgress', values: dict[str, Any]): + """Update the progress bar with new values. + + Args: + values (dict[str, Any]): The new values. + """ + for key, value in values.items(): + current_value = self.test_values.get(key, []) + self.test_values[key] = current_value + [value] class Trainer: @@ -136,8 +297,10 @@ class Trainer: epochs: int, optimizer: torch.optim.Optimizer, criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], - device: torch.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu'), val_loader: torch.utils.data.DataLoader | None = None, + test_loader: torch.utils.data.DataLoader | None = None, + metrics: List[Callable[[torch.Tensor, + torch.Tensor], torch.Tensor]] = [], ): """Train the model for the given number of epochs. @@ -147,12 +310,13 @@ class Trainer: epochs (int): The number of epochs to train the model for. optimizer (torch.optim.Optimizer): The optimizer to use. criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use. - device (torch.device, optional): The device to use. Defaults to torch.device('cuda' if torch.cuda.is_available() else 'cpu'). val_loader (torch.utils.data.DataLoader, optional): The validation dataset. Defaults to None. """ with TrainProgress( nb_epochs=epochs, - epoch_size=len(train_loader), + train_size=len(train_loader), + val_size=len(val_loader) if val_loader else 0, + test_size=len(test_loader) if test_loader else 0, ) as progress: self.progress = progress @@ -162,15 +326,20 @@ class Trainer: train_loader, optimizer, criterion, - device, + metrics, ) if val_loader: self.validate( model, val_loader, - criterion, - device, + metrics + [criterion], ) + if test_loader: + self.test( + model, + test_loader, + metrics + [criterion], + ) def train_epoch( self: 'Trainer', @@ -178,7 +347,7 @@ class Trainer: train_loader: torch.utils.data.DataLoader, optimizer: torch.optim.Optimizer, criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], - device: torch.device, + metrics: list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]], ): """Train the model for one epoch. @@ -187,13 +356,10 @@ class Trainer: train_loader (torch.utils.data.DataLoader): The training dataset. optimizer (torch.optim.Optimizer): The optimizer to use. criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use. - device (torch.device): The device to use. + metrics (list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]]): The metrics to use. """ model.train() for batch in train_loader: - # Move the batch to the device. - batch = [b.to(device) for b in batch] - # Seprarate the inputs and labels. inputs = batch[:-1] labels = batch[-1] @@ -206,27 +372,66 @@ class Trainer: optimizer.step() # Update the progress bar. - self.progress.new_values(loss=loss.item()) + values = {metric.__name__: metric(output, labels) + for metric in metrics} + values[criterion.__name__] = loss.item() self.progress.step() + self.progress.new_train_values(values) def validate( self: 'Trainer', model: torch.nn.Module, val_loader: torch.utils.data.DataLoader, - criterion: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], - device: torch.device, + metrics: list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]], ): """Validate the model on the given validation dataset. Args: model (torch.nn.Module): The model to validate. val_loader (torch.utils.data.DataLoader): The validation dataset. - criterion (Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): The loss function to use. - device (torch.device): The device to use. + mectrics (list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]]): The metrics to use. + """ + model.eval() + with torch.no_grad(): + metrics_sum = {f'val_{metric.__name__}': 0 for metric in metrics} + for b_i, batch in enumerate(val_loader): + inputs = batch[:-1] + labels = batch[-1] + output = model(*inputs) + values = {f'val_{metric.__name__}': metric(output, labels) + for metric in metrics} + for key, value in values.items(): + metrics_sum[key] += value.item() + self.progress.step() + self.progress.new_val_values({ + key: value / (b_i + 1) for key, value in metrics_sum.items() + }) + + def test( + self: 'Trainer', + model: torch.nn.Module, + test_loader: torch.utils.data.DataLoader, + metrics: list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]], + ): + """Test the model on the given test dataset. + + Args: + model (torch.nn.Module): The model to test. + test_loader (torch.utils.data.DataLoader): The test dataset. + mectrics (list[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]]): The metrics to use. """ model.eval() with torch.no_grad(): - for batch in val_loader: - batch = [b.to(device) for b in batch] - output = model(batch[:-1]) - loss = criterion(output, batch[-1]) + metrics_sum = {f'test_{metric.__name__}': 0 for metric in metrics} + for b_i, batch in enumerate(test_loader): + inputs = batch[:-1] + labels = batch[-1] + output = model(*inputs) + values = {f'test_{metric.__name__}': metric(output, labels) + for metric in metrics} + for key, value in values.items(): + metrics_sum[key] += value.item() + self.progress.step() + self.progress.new_test_values({ + key: value / (b_i + 1) for key, value in metrics_sum.items() + }) |
