checked GAN code

Author: aladdinpersson

ML/Pytorch/GANs/1. SimpleGAN/fc_gan.py

@@ -1,3 +1,12 @@
+""" 
+Simple GAN using fully connected layers
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+* 2020-11-01: Initial coding
+* 2022-12-20: Small revision of code, checked that it works with latest PyTorch version
+"""
+
+
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -48,7 +57,10 @@ def forward(self, x):
 gen = Generator(z_dim, image_dim).to(device)
 fixed_noise = torch.randn((batch_size, z_dim)).to(device)
 transforms = transforms.Compose(
-    [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)),]
+    [
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,), (0.5,)),
+    ]
 )
 
 dataset = datasets.MNIST(root="dataset/", transform=transforms, download=True)
@@ -104,4 +116,4 @@ def forward(self, x):
                 writer_real.add_image(
                     "Mnist Real Images", img_grid_real, global_step=step
                 )
-                step += 1
+                step += 1

ML/Pytorch/GANs/2. DCGAN/model.py

@@ -1,5 +1,9 @@
 """
 Discriminator and Generator implementation from DCGAN paper
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+* 2020-11-01: Initial coding
+* 2022-12-20: Small revision of code, checked that it works with latest PyTorch version
 """
 
 import torch
@@ -11,9 +15,7 @@ def __init__(self, channels_img, features_d):
         super(Discriminator, self).__init__()
         self.disc = nn.Sequential(
             # input: N x channels_img x 64 x 64
-            nn.Conv2d(
-                channels_img, features_d, kernel_size=4, stride=2, padding=1
-            ),
+            nn.Conv2d(channels_img, features_d, kernel_size=4, stride=2, padding=1),
             nn.LeakyReLU(0.2),
             # _block(in_channels, out_channels, kernel_size, stride, padding)
             self._block(features_d, features_d * 2, 4, 2, 1),
@@ -34,7 +36,7 @@ def _block(self, in_channels, out_channels, kernel_size, stride, padding):
                 padding,
                 bias=False,
             ),
-            #nn.BatchNorm2d(out_channels),
+            # nn.BatchNorm2d(out_channels),
             nn.LeakyReLU(0.2),
         )
 
@@ -68,7 +70,7 @@ def _block(self, in_channels, out_channels, kernel_size, stride, padding):
                 padding,
                 bias=False,
             ),
-            #nn.BatchNorm2d(out_channels),
+            # nn.BatchNorm2d(out_channels),
             nn.ReLU(),
         )
 
@@ -82,6 +84,7 @@ def initialize_weights(model):
         if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):
             nn.init.normal_(m.weight.data, 0.0, 0.02)
 
+
 def test():
     N, in_channels, H, W = 8, 3, 64, 64
     noise_dim = 100
@@ -91,6 +94,8 @@ def test():
     gen = Generator(noise_dim, in_channels, 8)
     z = torch.randn((N, noise_dim, 1, 1))
     assert gen(z).shape == (N, in_channels, H, W), "Generator test failed"
+    print("Success, tests passed!")
 
 
-# test()
+if __name__ == "__main__":
+    test()

ML/Pytorch/GANs/2. DCGAN/train.py

@@ -1,6 +1,10 @@
 """
 Training of DCGAN network on MNIST dataset with Discriminator
 and Generator imported from models.py
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+* 2020-11-01: Initial coding
+* 2022-12-20: Small revision of code, checked that it works with latest PyTorch version
 """
 
 import torch
@@ -35,11 +39,12 @@
 )
 
 # If you train on MNIST, remember to set channels_img to 1
-dataset = datasets.MNIST(root="dataset/", train=True, transform=transforms,
-                       download=True)
+dataset = datasets.MNIST(
+    root="dataset/", train=True, transform=transforms, download=True
+)
 
 # comment mnist above and uncomment below if train on CelebA
-#dataset = datasets.ImageFolder(root="celeb_dataset", transform=transforms)
+# dataset = datasets.ImageFolder(root="celeb_dataset", transform=transforms)
 dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
 gen = Generator(NOISE_DIM, CHANNELS_IMG, FEATURES_GEN).to(device)
 disc = Discriminator(CHANNELS_IMG, FEATURES_DISC).to(device)
@@ -92,14 +97,10 @@
             with torch.no_grad():
                 fake = gen(fixed_noise)
                 # take out (up to) 32 examples
-                img_grid_real = torchvision.utils.make_grid(
-                    real[:32], normalize=True
-                )
-                img_grid_fake = torchvision.utils.make_grid(
-                    fake[:32], normalize=True
-                )
+                img_grid_real = torchvision.utils.make_grid(real[:32], normalize=True)
+                img_grid_fake = torchvision.utils.make_grid(fake[:32], normalize=True)
 
                 writer_real.add_image("Real", img_grid_real, global_step=step)
                 writer_fake.add_image("Fake", img_grid_fake, global_step=step)
 
-            step += 1
+            step += 1

ML/Pytorch/GANs/3. WGAN/model.py

@@ -2,6 +2,10 @@
 Discriminator and Generator implementation from DCGAN paper,
 with removed Sigmoid() as output from Discriminator (and therefor
 it should be called critic)
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+* 2020-11-01: Initial coding
+* 2022-12-20: Small revision of code, checked that it works with latest PyTorch version
 """
 
 import torch
@@ -93,6 +97,7 @@ def test():
     gen = Generator(noise_dim, in_channels, 8)
     z = torch.randn((N, noise_dim, 1, 1))
     assert gen(z).shape == (N, in_channels, H, W), "Generator test failed"
+    print("Success, tests passed!")
 
-
-# test()
+if __name__ == "__main__":
+    test()

ML/Pytorch/GANs/3. WGAN/train.py

@@ -1,5 +1,9 @@
 """
 Training of DCGAN network with WGAN loss
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+* 2020-11-01: Initial coding
+* 2022-12-20: Small revision of code, checked that it works with latest PyTorch version
 """
 
 import torch
@@ -9,6 +13,7 @@
 import torchvision.datasets as datasets
 import torchvision.transforms as transforms
 from torch.utils.data import DataLoader
+from tqdm import tqdm
 from torch.utils.tensorboard import SummaryWriter
 from model import Discriminator, Generator, initialize_weights
 
@@ -61,7 +66,7 @@
 
 for epoch in range(NUM_EPOCHS):
     # Target labels not needed! <3 unsupervised
-    for batch_idx, (data, _) in enumerate(loader):
+    for batch_idx, (data, _) in enumerate(tqdm(loader)):
         data = data.to(device)
         cur_batch_size = data.shape[0]
 
@@ -111,4 +116,4 @@
 
             step += 1
             gen.train()
-            critic.train()
+            critic.train()

ML/Pytorch/GANs/4. WGAN-GP/model.py

@@ -1,5 +1,9 @@
 """
 Discriminator and Generator implementation from DCGAN paper
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+* 2020-11-01: Initial coding
+* 2022-12-20: Small revision of code, checked that it works with latest PyTorch version
 """
 
 import torch
@@ -24,7 +28,12 @@ def __init__(self, channels_img, features_d):
     def _block(self, in_channels, out_channels, kernel_size, stride, padding):
         return nn.Sequential(
             nn.Conv2d(
-                in_channels, out_channels, kernel_size, stride, padding, bias=False,
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                bias=False,
             ),
             nn.InstanceNorm2d(out_channels, affine=True),
             nn.LeakyReLU(0.2),
@@ -53,7 +62,12 @@ def __init__(self, channels_noise, channels_img, features_g):
     def _block(self, in_channels, out_channels, kernel_size, stride, padding):
         return nn.Sequential(
             nn.ConvTranspose2d(
-                in_channels, out_channels, kernel_size, stride, padding, bias=False,
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride,
+                padding,
+                bias=False,
             ),
             nn.BatchNorm2d(out_channels),
             nn.ReLU(),

ML/Pytorch/GANs/4. WGAN-GP/train.py

@@ -1,5 +1,9 @@
 """
 Training of WGAN-GP
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+* 2020-11-01: Initial coding
+* 2022-12-20: Small revision of code, checked that it works with latest PyTorch version
 """
 
 import torch
@@ -10,6 +14,7 @@
 import torchvision.transforms as transforms
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
 from utils import gradient_penalty, save_checkpoint, load_checkpoint
 from model import Discriminator, Generator, initialize_weights
 
@@ -31,13 +36,14 @@
         transforms.Resize(IMAGE_SIZE),
         transforms.ToTensor(),
         transforms.Normalize(
-            [0.5 for _ in range(CHANNELS_IMG)], [0.5 for _ in range(CHANNELS_IMG)]),
+            [0.5 for _ in range(CHANNELS_IMG)], [0.5 for _ in range(CHANNELS_IMG)]
+        ),
     ]
 )
 
 dataset = datasets.MNIST(root="dataset/", transform=transforms, download=True)
 # comment mnist above and uncomment below for training on CelebA
-#dataset = datasets.ImageFolder(root="celeb_dataset", transform=transforms)
+# dataset = datasets.ImageFolder(root="celeb_dataset", transform=transforms)
 loader = DataLoader(
     dataset,
     batch_size=BATCH_SIZE,
@@ -66,7 +72,7 @@
 
 for epoch in range(NUM_EPOCHS):
     # Target labels not needed! <3 unsupervised
-    for batch_idx, (real, _) in enumerate(loader):
+    for batch_idx, (real, _) in enumerate(tqdm(loader)):
         real = real.to(device)
         cur_batch_size = real.shape[0]
 
@@ -108,4 +114,4 @@
                 writer_real.add_image("Real", img_grid_real, global_step=step)
                 writer_fake.add_image("Fake", img_grid_fake, global_step=step)
 
-            step += 1
+            step += 1

ML/Pytorch/GANs/CycleGAN/config.py

@@ -11,7 +11,7 @@
 LAMBDA_CYCLE = 10
 NUM_WORKERS = 4
 NUM_EPOCHS = 10
-LOAD_MODEL = True
+LOAD_MODEL = False
 SAVE_MODEL = True
 CHECKPOINT_GEN_H = "genh.pth.tar"
 CHECKPOINT_GEN_Z = "genz.pth.tar"
@@ -24,6 +24,6 @@
         A.HorizontalFlip(p=0.5),
         A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255),
         ToTensorV2(),
-     ],
+    ],
     additional_targets={"image0": "image"},
-)
+)

ML/Pytorch/GANs/CycleGAN/discriminator_model.py

@@ -1,11 +1,28 @@
+"""
+Discriminator model for CycleGAN
+
+Programmed by Aladdin Persson <aladdin.persson at hotmail dot com>
+* 2020-11-05: Initial coding
+* 2022-12-21: Small revision of code, checked that it works with latest PyTorch version
+"""
+
 import torch
 import torch.nn as nn
 
+
 class Block(nn.Module):
     def __init__(self, in_channels, out_channels, stride):
         super().__init__()
         self.conv = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, 4, stride, 1, bias=True, padding_mode="reflect"),
+            nn.Conv2d(
+                in_channels,
+                out_channels,
+                4,
+                stride,
+                1,
+                bias=True,
+                padding_mode="reflect",
+            ),
             nn.InstanceNorm2d(out_channels),
             nn.LeakyReLU(0.2, inplace=True),
         )
@@ -32,15 +49,27 @@ def __init__(self, in_channels=3, features=[64, 128, 256, 512]):
         layers = []
         in_channels = features[0]
         for feature in features[1:]:
-            layers.append(Block(in_channels, feature, stride=1 if feature==features[-1] else 2))
+            layers.append(
+                Block(in_channels, feature, stride=1 if feature == features[-1] else 2)
+            )
             in_channels = feature
-        layers.append(nn.Conv2d(in_channels, 1, kernel_size=4, stride=1, padding=1, padding_mode="reflect"))
+        layers.append(
+            nn.Conv2d(
+                in_channels,
+                1,
+                kernel_size=4,
+                stride=1,
+                padding=1,
+                padding_mode="reflect",
+            )
+        )
         self.model = nn.Sequential(*layers)
 
     def forward(self, x):
         x = self.initial(x)
         return torch.sigmoid(self.model(x))
 
+
 def test():
     x = torch.randn((5, 3, 256, 256))
     model = Discriminator(in_channels=3)
@@ -50,4 +79,3 @@ def test():
 
 if __name__ == "__main__":
     test()
-