lfw cropped

1 files changed, 46 insertions, 33 deletions

diff --git a/main.py b/main.py
index ac11bab..b593775 100644
--- a/main.py
+++ b/main.py
@@ -32,19 +32,25 @@ class UNet(nn.Module):
         ## 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)  # 28x28 -> 14x14
+        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) # 14x14 -> 7x7
+        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) # 7x7 -> 14x14
-        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) # 14x14 -> 28x28
-        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)
+        self.maxpool3 = nn.MaxPool2d(2, 2)
+        self.conv7 = nn.Conv2d(256, 512, 3, padding=1)
+        self.conv8 = nn.Conv2d(512, 512, 3, padding=1)
+        self.upconv1 = nn.ConvTranspose2d(512, 256, 2, stride=2)
+        self.conv9 = nn.Conv2d(512, 256, 3, padding=1)
+        self.conv10 = nn.Conv2d(256, 256, 3, padding=1)
+        self.upconv2 = nn.ConvTranspose2d(256, 128, 2, stride=2)
+        self.conv11 = nn.Conv2d(256, 128, 3, padding=1)
+        self.conv12 = nn.Conv2d(128, 128, 3, padding=1)
+        self.upconv3 = nn.ConvTranspose2d(128, 64, 2, stride=2)
+        self.conv13 = nn.Conv2d(128, 64, 3, padding=1)
+        self.conv14 = nn.Conv2d(64, 64, 3, padding=1)
+        self.conv15 = nn.Conv2d(64, NB_CHANNEL, 3, padding=1)
 
     def forward(self, xt, t, vec):
         # Encode t and vec
@@ -65,17 +71,24 @@ class UNet(nn.Module):
         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 = self.maxpool3(x3)
         x4 = F.relu(self.conv7(x4))
         x4 = F.relu(self.conv8(x4))
-        x5 = self.upconv2(x4)
-        x5 = torch.cat((x5, x1), dim=1)
+        x5 = self.upconv1(x4)
+        x5 = torch.cat((x5, x3), dim=1)
         x5 = F.relu(self.conv9(x5))
         x5 = F.relu(self.conv10(x5))
-        x5 = self.conv11(x5)
+        x6 = self.upconv2(x5)
+        x6 = torch.cat((x6, x2), dim=1)
+        x6 = F.relu(self.conv11(x6))
+        x6 = F.relu(self.conv12(x6))
+        x7 = self.upconv3(x6)
+        x7 = torch.cat((x7, x1), dim=1)
+        x7 = F.relu(self.conv13(x7))
+        x7 = F.relu(self.conv14(x7))
+        x7 = self.conv15(x7)
 
-        return x5
+        return x7
 
 
 class LFWcrop(Dataset):
@@ -96,27 +109,27 @@ class LFWcrop(Dataset):
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-DIFFU_STEPS = 300
+DIFFU_STEPS = 1000
 BETA = torch.linspace(1e-4, 2e-2, DIFFU_STEPS, device=DEVICE)
 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.LFWPeople(
-    root="./data",
-    download=True,
-    transform=transforms.Compose([
-        transforms.Resize((64, 64)),
-        transforms.ToTensor(),
-    ]),
-)
+# dataset = datasets.MNIST(
+#     root="./data",
+#     train=True,
+#     download=True,
+#     transform=transforms.ToTensor(),
+# )
+# dataset = datasets.LFWPeople(
+#     root="./data",
+#     download=True,
+#     transform=transforms.Compose([
+#         transforms.Resize((64, 64)),
+#         transforms.ToTensor(),
+#     ]),
+# )
 dataset = LFWcrop()
 
 img = dataset[0][0]
@@ -159,7 +172,7 @@ def p_xt_1_xt(model, xt, t, vec):
     if sigma_theta.abs().max() <= 0:
         return mu_theta
 
-    return torch.distributions.Normal(mu_theta, sigma_theta).sample()
+    return torch.distributions.Normal(mu_theta, torch.sqrt(sigma_theta)).sample()
 
 
 class DiffusionDataset(Dataset):
@@ -244,7 +257,7 @@ if __name__ == '__main__':
                               num_workers=4, persistent_workers=True)
     trainer = Trainer()
     criterion = loss
-    epochs = 30
+    epochs = 100
 
     # Train the model.
     trainer.train(model, train_loader, epochs, optimizer, criterion)
@@ -260,7 +273,7 @@ if __name__ == '__main__':
     img, label = dataset[np.random.randint(0, len(dataset))]
     img = img.to(DEVICE) * 2 - 1
 
-    nb_plots = 6
+    nb_plots = 10
     plots_id = [i for i in np.linspace(1, DIFFU_STEPS, nb_plots, dtype=int)]
 
     xs = forward_diffusion(img)