Code

Network Architecture

  1. ### Train of DCGAN network on MNIST dataset with Discriminator
  2. ### and Generator imported from models.py
  5. import torch
  6. import torch.nn as nn
  8. #  they did not use batchNorm in the early in the first layer of the discriminator
  9. class Discriminator(nn.Module):
  10.     def __init__(self,channels_img, features_d):
  11.         super(Discriminator,self).__init__()
  12.         self.disc = nn.Sequential(
  13.             nn.Conv2d(channels_img, features_d, kernel_size=4,stride=2,padding=1),
  14.             nn.LeakyReLU(0.2),
  15.             self._block(features_d, features_d*2, 4, 2, 1),
  16.             self._block(features_d*2, features_d*4, 4, 2, 1),
  17.             self._block(features_d*4, features_d*8, 4, 2, 1),
  18.             nn.Conv2d(features_d*8, 1, kernel_size=4, stride=2, padding=0),
  19.             nn.Sigmoid(),
  21.         )
  24.     def _block(self, in_channels, out_channels, kernel_size, stride, padding):
  25.         return nn.Sequential(
  26.             nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=False),
  27.             nn.BatchNorm2d(out_channels),
  28.             nn.LeakyReLU(0.2),
  29.         )
  32.     def forward(self,x):
  33.         return self.disc(x)
  36. class Generator(nn.Module):
  37.     def __init__(self, z_dim, channels_img, features_g):
  38.         super(Generator,self).__init__()
  39.         self.gen = nn.Sequential(
  40.             # Input: N x z_dim x 1 x 1
  41.             self._block(z_dim, features_g*16, 4, 1, 0), # N x f_g*16 x 4 x 4
  42.             self._block(features_g*16, features_g*8, 4, 2 ,1),
  43.             self._block(features_g*8, features_g*4, 4, 2, 1),
  44.             self._block(features_g*4, features_g*2, 4, 2, 1),
  45.             nn.ConvTranspose2d(features_g*2, channels_img, kernel_size=4, stride=2, padding=1),
  46.             nn.Tanh(), #[-1,1]
  47.         )
  49.     def _block(self, in_channels, out_channels, kernel_size, stride, padding):
  50.         return nn.Sequential(
  51.             nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding, bias=False),
  52.             nn.BatchNorm2d(out_channels),
  53.             nn.ReLU(),
  54.         )
  56.     def forward(self,x):
  57.         return self.gen(x)
  60. def initialize_weights(model):
  61.     for m in model.modules():
  62.         if isinstance(m,(nn.Conv2d, nn.ConvTranspose2d, nn.BatchNorm2d)):
  63.             nn.init.normal_(m.weight.data,0.0,0.02)
  65. def test():
  66.     N, in_channels, H, W = 8, 3, 64, 64
  67.     z_dim = 100
  68.     x = torch.randn((N,in_channels, H, W))
  69.     disc = Discriminator(in_channels, 8)
  70.     initialize_weights(disc)
  71.     assert disc(x).shape == (N, 1, 1, 1)
  73.     gen = Generator(z_dim, in_channels, 8)
  74.     initialize_weights(gen)
  75.     z = torch.randn((N, z_dim, 1, 1))
  76.     assert gen(z).shape == (N, in_channels, H, W)
  77.     print("Success!")

Training

  2. import torch
  3. import torch.nn as nn
  4. import torch.optim as optim
  5. import torchvision
  6. import torchvision.datasets as datasets
  7. import torchvision.transforms as transforms
  8. from torch.utils.data import DataLoader
  9. from torch.utils.tensorboard import SummaryWriter
  10. from model import Discriminator, Generator, initialize_weights
  13. ###  Hyper parameters etc.
  14. device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  15. print(torch.cuda.is_available())
  16. Learning_rate = 2e-4
  17. Batch_size = 128
  18. Image_size = 64
  19. CHANNELS_IMG = 1
  20. Noise_dim = 100
  21. Num_epochs = 5
  22. Features_disc = 64
  23. Features_gen = 64
  25. ### Image pre-processing
  26. transformx = transforms.Compose(
  27.     [
  28.         transforms.Resize(Image_size),
  29.         transforms.ToTensor(),
  30.         transforms.Normalize(
  31.             [0.5 for _ in range(CHANNELS_IMG)],[0.5 for _ in range(CHANNELS_IMG)]
  32.         ),
  33.     ]
  34. )
  36. ### Load dataset
  37. dataset = datasets.MNIST(root="./data",train=True,transform=transformx,download=False)
  38. loader = DataLoader(dataset, batch_size=Batch_size, shuffle=True)
  41. ### Build the model
  43. gen = Generator(Noise_dim, CHANNELS_IMG, Features_gen).to(device)
  44. disc = Discriminator(CHANNELS_IMG,Features_disc).to(device)
  46. # initialize model parameter
  47. initialize_weights(gen)
  48. initialize_weights(disc)
  50. opt_gen = optim.Adam(gen.parameters(),lr=Learning_rate, betas=(0.5, 0.999))
  51. opt_disc = optim.Adam(disc.parameters(),lr=Learning_rate,betas=(0.5, 0.999))
  52. # loss function to calculate loss = y_n*log(x_n) + (1-y_n)*log(1-x_n)
  53. criterion = nn.BCELoss()
  55. # the fixed_noise is used to validate the model
  56. fixed_noise = torch.randn(32,Noise_dim,1,1).to(device)
  58. step = 0 # need a step for printing to tensorboard
  59. writer_real = SummaryWriter(f"logs/real")
  60. writer_fake = SummaryWriter(f"logs/fake")
  63. gen.train()
  64. disc.train()
  66. ## start training
  68. for epoch in range(Num_epochs):
  69.     for batch_idx, (real, _) in enumerate(loader):
  70.         # get real image from training set
  71.         real = real.to(device)
  72.         # generate fake image
  73.         noise = torch.randn((Batch_size,Noise_dim,1,1)).to(device)
  74.         fake = gen(noise)
  75.         # Train discriminator max log(D(x)) + log(1-D(G(z)))
  76.         disc_real = disc(real).reshape(-1)
  77.         loss_disc_real = criterion(disc_real, torch.ones_like(disc_real))
  78.         disc_fake = disc(fake).reshape(-1)
  79.         loss_disc_fake = criterion(disc_fake, torch.zeros_like(disc_fake))
  80.         # generate discriminator loss
  81.         loss_disc = (loss_disc_real + loss_disc_fake)/2
  82.         disc.zero_grad()
  83.         loss_disc.backward(retain_graph=True)
  84.         opt_disc.step()
  86.         ### Train the Generator min log(1-D(G(z)) <--> max log(D(G(z)))
  87.         output = disc(fake).reshape(-1)
  88.         loss_gen = criterion(output,torch.ones_like(output))
  89.         gen.zero_grad()
  90.         loss_gen.backward()
  91.         opt_gen.step()
  94.         # Print losses occasionally and print to tensorboard
  95.         if batch_idx % 100 == 0:
  96.             print(
  97.                 f"Epoch [{epoch}/{Num_epochs}] Batch {batch_idx}/{len(loader)} \
  98.                           Loss D: {loss_disc:.4f}, loss G: {loss_gen:.4f}"
  99.             )
  101.             with torch.no_grad():
  102.                 fake = gen(fixed_noise)
  103.                 # take out (up to) 32 examples
  104.                 img_grid_real = torchvision.utils.make_grid(
  105.                     real[:32], normalize=True
  106.                 )
  107.                 img_grid_fake = torchvision.utils.make_grid(
  108.                     fake[:32], normalize=True
  109.                 )
  111.                 writer_real.add_image("Real", img_grid_real, global_step=step)
  112.                 writer_fake.add_image("Fake", img_grid_fake, global_step=step)
  114.             step += 1
  117. import matplotlib.pyplot as plt
  119. plt.figure(figsize=(10,5))
  120. plt.title("Generator and Discriminator Loss During Training")
  121. plt.plot(loss_gen.data.cpu().numpy(),label="G")
  122. plt.plot(loss_disc.data.cpu().numpy(),label="D")
  123. plt.xlabel("iterations")
  124. plt.ylabel("Loss")
  125. plt.legend()
  126. plt.show()