diff options
| -rw-r--r-- | main.py | 250 |
1 files changed, 169 insertions, 81 deletions
@@ -37,7 +37,7 @@ class SelfAttention(nn.Module): return x -class UNetEDF(nn.Module): +class UNetAddAttUEDF(nn.Module): def __init__(self): super().__init__() @@ -57,15 +57,15 @@ class UNetEDF(nn.Module): self.conv2 = nn.Conv2d(64, 64, 3, padding=1) self.maxpool1 = nn.MaxPool2d(2, 2) self.conv3 = nn.Conv2d(64, 128, 3, padding=1) - # self.att1 = SelfAttention(128, 8) self.conv4 = nn.Conv2d(128, 128, 3, padding=1) self.maxpool2 = nn.MaxPool2d(2, 2) + self.att1 = SelfAttention(128, 8) self.conv5 = nn.Conv2d(128, 256, 3, padding=1) - # self.att2 = SelfAttention(256, 8) + self.att2 = SelfAttention(256, 8) self.conv6 = nn.Conv2d(256, 256, 3, padding=1) + self.att3 = SelfAttention(256, 8) self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2) self.conv7 = nn.Conv2d(256, 128, 3, padding=1) - # self.att3 = SelfAttention(128, 8) self.conv8 = nn.Conv2d(128, 128, 3, padding=1) self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2) self.conv9 = nn.Conv2d(128, 64, 3, padding=1) @@ -87,24 +87,87 @@ class UNetEDF(nn.Module): x1 = F.relu(self.conv2(x1)) x2 = self.maxpool1(x1) x2 = F.relu(self.conv3(x2)) - # x2 = self.att1(x2) x2 = F.relu(self.conv4(x2)) x3 = self.maxpool2(x2) + x3 = self.att1(x3) x3 = F.relu(self.conv5(x3)) - # x3 = self.att2(x3) - x5 = F.relu(self.conv6(x3)) - x6 = self.upconv1(x5) - x6 = torch.cat((x6, x2), dim=1) - x6 = F.relu(self.conv7(x6)) - # x6 = self.att3(x6) - x6 = F.relu(self.conv8(x6)) - x7 = self.upconv2(x6) - x7 = torch.cat((x7, x1), dim=1) - x7 = F.relu(self.conv9(x7)) - x7 = F.relu(self.conv10(x7)) - x7 = self.conv11(x7) + x3 = self.att2(x3) + x3 = F.relu(self.conv6(x3)) + x4 = self.upconv1(x3) + x4 = torch.cat((x4, x2), dim=1) + x4 = F.relu(self.conv7(x4)) + x4 = F.relu(self.conv8(x4)) + x5 = self.upconv2(x4) + x5 = torch.cat((x5, x1), dim=1) + x5 = F.relu(self.conv9(x5)) + x5 = F.relu(self.conv10(x5)) + x5 = self.conv11(x5) - return x7 + return x5 + + +class UNetMNISTSig(nn.Module): + def __init__(self): + super().__init__() + + ## Inputs: + # xt: image at step t (NB_CHANNEL*IMG_SIZE*IMG_SIZE) + # t: step number (1) + # vec: one-hot vector of the label (NB_LABEL) + + ## Encoder for t + self.encodet = nn.Linear(1, IMG_SIZE*IMG_SIZE) + + ## Encoder for vec + self.encodevec = nn.Linear(NB_LABEL, IMG_SIZE*IMG_SIZE) + + ## UNet (2 more channels input because we concatenate xt with t and vec) + self.conv1 = nn.Conv2d(NB_CHANNEL+2, 64, 3, padding=1) + self.conv2 = nn.Conv2d(64, 64, 3, padding=1) + self.maxpool1 = nn.MaxPool2d(2, 2) + self.conv3 = nn.Conv2d(64, 128, 3, padding=1) + self.conv4 = nn.Conv2d(128, 128, 3, padding=1) + self.maxpool2 = nn.MaxPool2d(2, 2) + self.conv5 = nn.Conv2d(128, 256, 3, padding=1) + self.conv6 = nn.Conv2d(256, 256, 3, padding=1) + self.upconv1 = nn.ConvTranspose2d(256, 128, 2, stride=2) + self.conv7 = nn.Conv2d(256, 128, 3, padding=1) + self.conv8 = nn.Conv2d(128, 128, 3, padding=1) + self.upconv2 = nn.ConvTranspose2d(128, 64, 2, stride=2) + self.conv9 = nn.Conv2d(128, 64, 3, padding=1) + self.conv10 = nn.Conv2d(64, 64, 3, padding=1) + self.conv11 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1) + + def forward(self, xt, t, vec): + # Encode t and vec + t = F.relu(self.encodet(t / DIFFU_STEPS)) + t = t.view(-1, 1, IMG_SIZE, IMG_SIZE) + vec = F.relu(self.encodevec(vec)) + vec = vec.view(-1, 1, IMG_SIZE, IMG_SIZE) + + # Concat all 3 + x = torch.cat((xt, t, vec), dim=1) + + # UNet + x1 = F.relu(self.conv1(x)) + x1 = F.relu(self.conv2(x1)) + x2 = self.maxpool1(x1) + x2 = F.relu(self.conv3(x2)) + x2 = F.relu(self.conv4(x2)) + x3 = self.maxpool2(x2) + x3 = F.relu(self.conv5(x3)) + x3 = F.relu(self.conv6(x3)) + x4 = self.upconv1(x3) + x4 = torch.cat((x4, x2), dim=1) + x4 = F.relu(self.conv7(x4)) + x4 = F.relu(self.conv8(x4)) + x5 = self.upconv2(x4) + x5 = torch.cat((x5, x1), dim=1) + x5 = F.relu(self.conv9(x5)) + x5 = F.relu(self.conv10(x5)) + x5 = self.conv11(x5) + + return x5 class FolderDataset(Dataset): def __init__(self, path, size=(32, 32)): @@ -136,9 +199,20 @@ def q_xt_xt_1(xt_1, t): return xt, eps + def q_xt_x0(x0, t): t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t + reshaped = len(x0.shape) == 3 + if reshaped: + c, w, h = x0.shape + b = 1 + x0 = x0.view(b, c, w, h) + else: + b, c, w, h = x0.shape + t_ind = t_ind.view(b, 1, 1, 1) + t_ind = t_ind.expand(b, c, w, h) + alpha_bar = ALPHA_BAR[t_ind] mean = torch.sqrt(alpha_bar) * x0 std = torch.sqrt(1 - alpha_bar) @@ -146,8 +220,12 @@ def q_xt_x0(x0, t): eps = torch.randn(x0.shape, device=DEVICE) xt = mean + std * eps + if reshaped: + xt = xt.view(c, w, h) + return xt, eps + def p_xt_1_xt(model, xt, t, vec): t_ind = t.to(dtype=torch.long) if isinstance(t, torch.Tensor) else t @@ -169,6 +247,12 @@ def p_xt_1_xt(model, xt, t, vec): return mu_theta + sigma_theta * noise +def p_xt_1_xt_sig(model, xt_1, t, vec): + x0 = model(xt_1, t, vec) + xt, _ = q_xt_x0(x0, t) + return xt + + class DiffusionDataset(Dataset): def __init__(self, dataset, autoencoder=None): super().__init__() @@ -210,6 +294,47 @@ class DiffusionDataset(Dataset): return len(self.dataset) +class DiffusionDatasetSig(Dataset): + def __init__(self, dataset, autoencoder=None): + super().__init__() + self.dataset = dataset + self.autoencoder = autoencoder + + def __getitem__(self, index): + # Get the image and the label. + img, label = self.dataset[index] + img = img.to(DEVICE) + + # Encode the image. + if self.autoencoder is not None: + img = self.autoencoder.encode(img.unsqueeze(0)).squeeze(0) + + # Normalize the image. + img = img * 2 - 1 + + # Add noise to the image. + t = torch.randint(1, DIFFU_STEPS, (1,), device=DEVICE) + xt, eps = q_xt_x0(img, t) + + # Convert the label to a one-hot vector. + vec = torch.nn.functional.one_hot( + torch.tensor(min(label, NB_LABEL-1)), + num_classes=NB_LABEL, + ) + + return ( + # x_true + xt.clone().detach().to(dtype=torch.float32, device=DEVICE), + t.clone().detach().to(dtype=torch.float32, device=DEVICE), + vec.clone().detach().to(dtype=torch.float32, device=DEVICE), + # y_true + img.clone().detach().to(dtype=torch.float32, device=DEVICE), + ) + + def __len__(self): + return len(self.dataset) + + def loss(y_pred, y_true): return nn.MSELoss()(y_pred, y_true) @@ -218,7 +343,7 @@ def forward_diffusion(x0): x = x0.clone() xs = [x] for t in range(1, DIFFU_STEPS+1): - x = q_xt_xt_1(x, t)[0] + x, _ = q_xt_xt_1(x, t) xs.append(x) return xs @@ -232,8 +357,8 @@ def tensor_to_image(tensor): def tensor_to_images(tensor): img = tensor.clone().detach().cpu().numpy().transpose(0, 2, 3, 1) - img -= img.min() - img /= img.max() + img -= np.min(img, axis=(1, 2, 3), keepdims=True) + img /= np.max(img, axis=(1, 2, 3), keepdims=True) return img @@ -253,12 +378,12 @@ BETA = torch.cat((torch.tensor([0.], device=DEVICE), BETA)) ALPHA = 1 - BETA ALPHA_BAR = torch.cumprod(ALPHA, dim=0) -# dataset = datasets.MNIST( -# root="./data", -# train=True, -# download=True, -# transform=transforms.ToTensor(), -# ) +dataset = datasets.MNIST( + root="./data", + train=True, + download=True, + transform=transforms.ToTensor(), +) # dataset = datasets.LFWPeople( # root="./data", # download=True, @@ -268,7 +393,7 @@ ALPHA_BAR = torch.cumprod(ALPHA, dim=0) # ]), # ) # dataset = FolderDataset('data/lfwcrop_color/faces') -dataset = FolderDataset('data/edface') +# dataset = FolderDataset('data/edface') autoencoder = None # autoencoder = Autoencoder(1, 1).to(DEVICE) @@ -279,16 +404,16 @@ img = dataset[0][0].to(DEVICE) if autoencoder is not None: img = autoencoder.encode(img.unsqueeze(0)).squeeze(0) NB_CHANNEL, IMG_SIZE, _ = img.shape -NB_LABEL = 1 +NB_LABEL = 10 -EPOCHS = 200 +EPOCHS = 10 LEARNING_RATE = 2e-4 def epoch_callback(trainer: Trainer): epoch_i = trainer.epoch_i - if epoch_i % 10 == 0 or epoch_i == trainer.epoch_end: + if epoch_i % 1 == 0 or epoch_i == trainer.epoch_end: save_path = f'runs/{trainer.run_name}/checkpoints/{epoch_i:04}e.pt' torch.save(trainer.model, save_path) save_path = f'runs/{trainer.run_name}/checkpoints/last.pt' @@ -308,7 +433,7 @@ def epoch_callback(trainer: Trainer): for t in range(DIFFU_STEPS, 0, -1): t_tensor = torch.tensor([[t]] * batch_size, device=DEVICE, dtype=torch.float32) - x = p_xt_1_xt(model, x, t_tensor, vec) + x = p_xt_1_xt_sig(model, x, t_tensor, vec) x = x.cpu().numpy() x -= x.min(axis=(1, 2, 3), keepdims=True) @@ -324,9 +449,9 @@ def epoch_callback(trainer: Trainer): trainer.writer.add_scalar('JSD/Validation', jsd(reals, fakes), epoch_i) - fig = plt.figure(figsize=(16, 16)) - for i in range(16): - plt.subplot(4, 4, i + 1) + fig = plt.figure(figsize=(32, 16)) + for i in range(32): + plt.subplot(4, 8, i + 1) plt.imshow(fakes[i].transpose(1, 2, 0)) plt.axis("off") plt.tight_layout() @@ -343,14 +468,14 @@ if __name__ == '__main__': torch.multiprocessing.set_start_method("spawn") # Load the model. - # model = UNetEDF().to(DEVICE) - model = torch.load('runs/UNetEDF_20240702-141225/checkpoints/last.pt') + # model = UNetMNISTSig().to(DEVICE) + model = torch.load('runs/20240719-154613_UNetMNISTSig/checkpoints/last.pt').to(DEVICE) # Define the optimizer. optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE) # Define the training dataset. - train_dataset = DiffusionDataset(dataset, autoencoder) + train_dataset = DiffusionDatasetSig(dataset, autoencoder) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=4, persistent_workers=True) criterion = loss @@ -385,7 +510,7 @@ if __name__ == '__main__': plt.axis("off") for plot_i, t in enumerate(plots_id): - x = q_xt_x0(img, t)[0] + x, _ = q_xt_x0(img, t) plt.subplot(2, nb_plots + 1, nb_plots + plot_i + 3) plt.imshow(tensor_to_image(x), interpolation='none') plt.axis("off") @@ -399,7 +524,7 @@ if __name__ == '__main__': plt.suptitle("Forward diffusion") plt.tight_layout() - plt.savefig("forward_diffusion.tmp.png") + plt.savefig("plots/forward_diffusion.tmp.png") # Backward diffusion @@ -415,7 +540,7 @@ if __name__ == '__main__': plt.suptitle("Backward diffusion") for t in track(range(DIFFU_STEPS, 0, -1), description='Diffusing...'): t_tensor = torch.tensor([[t]] * n_classes, device=DEVICE, dtype=torch.float32) - x = p_xt_1_xt(model, x, t_tensor, vec) + x = p_xt_1_xt_sig(model, x, t_tensor, vec) if t in t_plots: t_plot_i = nb_plots - t_plots.tolist().index(t) - 1 for class_i in range(n_classes): @@ -425,7 +550,7 @@ if __name__ == '__main__': plt.imshow(tensor_to_image(x[class_i])) plt.axis("off") plt.tight_layout() - plt.savefig("backward_diffusion.tmp.png") + plt.savefig("plots/backward_diffusion.tmp.png") # Benchmark @@ -436,7 +561,7 @@ if __name__ == '__main__': for ti in track(range(DIFFU_STEPS, 0, -1), description='Benchmarking...'): t = torch.tensor([[ti]] * n_classes * nb_plots, device=DEVICE, dtype=torch.float32) - x = p_xt_1_xt(model, x, t, vec) + x = p_xt_1_xt_sig(model, x, t, vec) x = x * 0.5 + 0.5 x = x.clamp(0, 1) @@ -452,41 +577,4 @@ if __name__ == '__main__': plt.imshow(tensor_to_image(img)) plt.axis("off") plt.tight_layout() - plt.savefig("benchmark.tmp.png") - - - # # Metrics - # batch_size = 64 - # n_batches = 16 - # n_samples = batch_size * n_batches - - # fakes = np.zeros((0, NB_CHANNEL, IMG_SIZE, IMG_SIZE)) - # for _ in track(range(n_batches), description=f'Sampling {n_samples} images...'): - # x = torch.randn(batch_size, NB_CHANNEL, IMG_SIZE, IMG_SIZE).to(DEVICE) - # vec = torch.randint(0, NB_LABEL, (batch_size,)).to(DEVICE) - # vec = torch.nn.functional.one_hot(vec, num_classes=NB_LABEL).to(device=DEVICE, dtype=torch.float32) - - # for t in range(DIFFU_STEPS, 0, -1): - # t_tensor = torch.tensor([[t]] * batch_size, device=DEVICE, dtype=torch.float32) - # x = p_xt_1_xt(model, x, t_tensor, vec) - - # fakes = np.concatenate((fakes, x.cpu().numpy())) - - # reals = torch.stack([dataset[i][0] for i in range(n_samples)]).cpu().numpy() - # reals = reals * 2 - 1 - - - # fid_score = fid(reals, fakes) - # print(f"FID score: {fid_score}") - - # kl_score = kl(reals, fakes) - # print(f"KL divergence: {kl_score}") - - # rkl_score = kl(fakes, reals) - # print(f"Reverse KL divergence: {rkl_score}") - - # jsd_score = jsd(reals, fakes) - # print(f"Jensen-Shannon divergence: {jsd_score}") - - - # plt.show() + plt.savefig("plots/benchmark.tmp.png") |