DCGAN initial commit

caffeine-coder1 · caffeine-coder1 · commit b3f6c5a119ce · 2021-06-18T12:05:42.000+05:30
diff --git a/GAN/DCGAN/model.py b/GAN/DCGAN/model.py
@@ -0,0 +1,77 @@
+import torch.nn as nn
+from torchsummary import summary
+
+
+class Discriminator(nn.Module):
+    def __init__(self, img_channels, feature_d):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(img_channels, feature_d,
+                      kernel_size=4, stride=2, padding=1),
+            nn.LeakyReLU(0.2),
+            *self.__block(feature_d, feature_d*2,
+                          kernel_size=4, stride=2, padding=1),
+            *self.__block(feature_d*2, feature_d*4,
+                          kernel_size=4, stride=2, padding=1),
+            *self.__block(feature_d*4, feature_d*8,
+                          kernel_size=4, stride=2, padding=1),
+            nn.Conv2d(feature_d*8, 1,
+                      kernel_size=4, stride=2, padding=0),
+            nn.Sigmoid()
+        )
+        self.initialize_weights()
+
+    def __block(self, in_channels, out_channels, **kwargs):
+        return [nn.Conv2d(in_channels, out_channels, **kwargs, bias=False),
+                nn.BatchNorm2d(out_channels), nn.LeakyReLU(0.2)]
+
+    def initialize_weights(self):
+        for m in self.net.modules():
+            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):
+                nn.init.normal_(m.weight.data, 0, 0.02)
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class Faker(nn.Module):
+
+    def __init__(self, z_dim, img_channels, feature_d):
+        super().__init__()
+        self.net = nn.Sequential(
+            *self.__block(z_dim, feature_d*8, kernel_size=4,
+                          stride=2, padding=0),
+            *self.__block(feature_d*8, feature_d*4, kernel_size=4,
+                          stride=2, padding=1),
+            *self.__block(feature_d*4, feature_d*2, kernel_size=4,
+                          stride=2, padding=1),
+            *self.__block(feature_d*2, feature_d, kernel_size=4,
+                          stride=2, padding=1),
+            nn.ConvTranspose2d(feature_d, img_channels,
+                               kernel_size=4, stride=2, padding=1),
+            nn.Tanh()
+        )
+        self.initialize_weights()
+
+    def __block(self, in_channels, out_channels, **kwargs):
+
+        return [nn.ConvTranspose2d(in_channels, out_channels, **kwargs, bias=False),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU()]
+
+    def initialize_weights(self):
+        for m in self.net.modules():
+            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):
+                nn.init.normal_(m.weight.data, 0, 0.02)
+
+    def forward(self, x):
+        return self.net(x)
+
+
+if __name__ == "__main__":
+    disc = Discriminator(3, 128)
+    gen = Faker(100, 3, 128)
+
+    summary(disc, (3, 64, 64))
+
+    summary(gen, (100, 1, 1))
diff --git a/GAN/DCGAN/training.py b/GAN/DCGAN/training.py
@@ -0,0 +1,117 @@
+import torch
+import torchvision
+import torch.optim as optim
+from torchvision.transforms import transforms
+from torchvision.datasets import MNIST
+from torch.utils.data import DataLoader
+from model import Discriminator, Faker
+from torch.utils.tensorboard import SummaryWriter
+
+
+# ~~~~~~~~~~~~~~~~~~~ hyper parameters ~~~~~~~~~~~~~~~~~~~ #
+EPOCHS = 20
+CHANNELS = 1
+H, W = 64, 64
+IMG_SIZE = CHANNELS * H * W
+lr = 2e-4
+work_device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+FEATURE_D = 128
+Z_DIM = 100
+GEN_TRAIN_STEPS = 5
+BATCH_SIZE = 128
+# ~~~~~~~~~~~~~~~~~~~ loading the dataset ~~~~~~~~~~~~~~~~~~~ #
+
+trans = transforms.Compose(
+    [transforms.Resize((H, W)), transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
+
+MNIST_data = MNIST('./data', True, transform=trans, download=True)
+
+loader = DataLoader(MNIST_data, BATCH_SIZE, True, num_workers=1)
+
+# ~~~~~~~~~~~~~~~~~~~ creating tensorboard variables ~~~~~~~~~~~~~~~~~~~ #
+
+writer_fake = SummaryWriter("logs/fake")
+writer_real = SummaryWriter("logs/real")
+
+# ~~~~~~~~~~~~~~~~~~~ loading the model ~~~~~~~~~~~~~~~~~~~ #
+
+disc = Discriminator(img_channels=CHANNELS,
+                     feature_d=FEATURE_D).to(work_device)
+gen = Faker(Z_DIM, CHANNELS, FEATURE_D).to(work_device)
+
+# ~~~~~~~~~~~~~~~~~~~ create optimizer and loss ~~~~~~~~~~~~~~~~~~~ #
+
+disc_optim = optim.Adam(disc.parameters(), lr, (0.5, 0.999))
+gen_optim = optim.Adam(gen.parameters(), lr, (0.5, 0.999))
+criterion = torch.nn.BCELoss()
+
+# ~~~~~~~~~~~~~~~~~~~ training loop ~~~~~~~~~~~~~~~~~~~ #
+
+for epoch in range(EPOCHS):
+
+    for batch_idx, (real, _) in enumerate(loader):
+        disc.train()
+        gen.train()
+        real = real.to(work_device)
+        fixed_noise = torch.rand(real.shape[0], Z_DIM, H, W).to(work_device)
+        # ~~~~~~~~~~~~~~~~~~~ discriminator loop ~~~~~~~~~~~~~~~~~~~ #
+
+        fake = gen(fixed_noise)  # dim of (N,1,28,28)
+        # ~~~~~~~~~~~~~~~~~~~ forward ~~~~~~~~~~~~~~~~~~~ #
+        real_predict = disc(real).view(-1)  # make it one dimensional array
+        fake_predict = disc(fake).view(-1)  # make it one dimensional array
+
+        labels = torch.cat([torch.ones_like(real_predict),
+                            torch.zeros_like(fake_predict)], dim=0)
+
+        # ~~~~~~~~~~~~~~~~~~~ loss ~~~~~~~~~~~~~~~~~~~ #
+        D_loss = criterion(
+            torch.cat([real_predict, fake_predict], dim=0), labels)
+
+        # ~~~~~~~~~~~~~~~~~~~ backward ~~~~~~~~~~~~~~~~~~~ #
+        disc.zero_grad()
+        D_loss.backward()
+        disc_optim.step()
+
+        # ~~~~~~~~~~~~~~~~~~~ generator loop ~~~~~~~~~~~~~~~~~~~ #
+        for _ in range(GEN_TRAIN_STEPS):
+            # ~~~~~~~~~~~~~~~~~~~ forward ~~~~~~~~~~~~~~~~~~~ #
+            fake = gen(fixed_noise).view(-1, CHANNELS,
+                                         H, W)  # dim of (N,1,32,32)
+            # ~~~~~~~~~~~~~~~~~~~ forward ~~~~~~~~~~~~~~~~~~~ #
+            fake_predict = disc(fake).view(-1)  # make it one dimensional array
+            # ~~~~~~~~~~~~~~~~~~~ loss ~~~~~~~~~~~~~~~~~~~ #
+
+            G_loss = criterion(fake_predict, torch.ones_like(fake_predict))
+            # ~~~~~~~~~~~~~~~~~~~ backward ~~~~~~~~~~~~~~~~~~~ #
+            gen.zero_grad()
+            G_loss.backward()
+            gen_optim.step()
+
+        # ~~~~~~~~~~~~~~~~~~~ loading the tensorboard ~~~~~~~~~~~~~~~~~~~ #
+
+        if batch_idx == 0:
+            print(
+                f"Epoch [{epoch}/{EPOCHS}] Batch {batch_idx}/{len(loader)} \
+                            Loss D: {D_loss:.4f}, loss G: {G_loss:.4f}"
+            )
+
+            with torch.no_grad():
+                disc.eval()
+                gen.eval()
+                fake = gen(fixed_noise).reshape(-1, CHANNELS, H, W)
+                data = real.reshape(-1, CHANNELS, H, W)
+                if BATCH_SIZE > 32:
+                    fake = fake[:32]
+                    data = data[:32]
+                img_grid_fake = torchvision.utils.make_grid(
+                    fake, normalize=True)
+                img_grid_real = torchvision.utils.make_grid(
+                    data, normalize=True)
+
+                writer_fake.add_image(
+                    "Mnist Fake Images", img_grid_fake, global_step=epoch
+                )
+                writer_real.add_image(
+                    "Mnist Real Images", img_grid_real, global_step=epoch
+                )
