idk

1 files changed, 169 insertions, 81 deletions

diff --git a/main.py b/main.py
index 11b1824..7f8444a 100644
--- a/main.py
+++ b/main.py
@@ -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")