Search⌘ K
Home
Courses
Hands-On Generative Adversarial Networks with PyTorch

Creating a DCGAN with PyTorch

Explore the step-by-step process of creating a Deep Convolutional GAN (DCGAN) with PyTorch. Learn how to define generator and discriminator networks, initialize model parameters, and implement training using Adam and Binary Cross-Entropy loss. Understand dataset handling, training iteration steps, and sample visualization to evaluate model performance on generating handwritten digit images.

We'll cover the following...

Let’s start writing PyTorch code to create a DCGAN model. Here, we assume that the Python environment is being used in Ubuntu.

First, let’s create a Python source file called dcgan.py and import the packages that we need:

Python 3.10.4
import os
import sys
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
import utils

Here, numpy is only used to initialize a random seed. If we don't have numpy installed, simply replace np.random with random and insert the import random line after import os. In the last line of code, we import a module called utils, which is a custom utility package defined in the utils.py file. The full source code of utils.py is as follows:

Python 3.10.4
import errno
import os
import shutil
import sys
def to_np(var):
    """Exports torch.Tensor to Numpy array.
    """
    return var.detach().cpu().numpy()
def create_folder(folder_path):
    """Create a folder if it does not exist.
    """
    try:
        os.makedirs(folder_path)
    except OSError as _e:
        if _e.errno != errno.EEXIST:
            raise
def clear_folder(folder_path):
    """Clear all contents recursively if the folder exists.
    Create the folder if it has been accidently deleted.
    """
    create_folder(folder_path)
    for the_file in os.listdir(folder_path):
        _file_path = os.path.join(folder_path, the_file)
        try:
            if os.path.isfile(_file_path):
                os.unlink(_file_path)
            elif os.path.isdir(_file_path):
                shutil.rmtree(_file_path)
        except OSError as _e:
            print(_e)
class StdOut(object):
    """Redirect stdout to file, and print to console as well.
    """
    def __init__(self, output_file):
        self.terminal = sys.stdout
        self.log = open(output_file, "a")
    def write(self, message):
        self.terminal.write(message)
        self.terminal.flush()
        self.log.write(message)
        self.log.flush()
    def flush(self):
        self.terminal.flush()
        self.log.flush()

We will put most of the PyTorch-independent helper functions (including file organization, learning rate adjustment, logging, tensor visualization, and so on) in this utils.py file.

Then, we define the output path and hyperparameters. Note that here we set the minimal channel size of hidden layers in both the generator and discriminator to 64:

Python 3.10.4
CUDA = True
DATA_PATH = '~/Data/mnist'
OUT_PATH = 'output'
LOG_FILE = os.path.join(OUT_PATH, 'log.txt')
BATCH_SIZE = 128
IMAGE_CHANNEL = 1
Z_DIM = 100
G_HIDDEN = 64
X_DIM = 64
D_HIDDEN = 64
EPOCH_NUM = 25
REAL_LABEL = 1
FAKE_LABEL = 0
lr = 2e-4
seed = 1

If we don't have a CUDA-enabled graphics card and want to train the networks on the CPU, we can change CUDA to False. DATA_PATH points to the root directory of the MNIST dataset. If MNIST hasn't been downloaded and properly preprocessed yet, simply point it to any directory (such as ...