How to Build a DCGAN with PyTorch | by Conor Lazarou | Towards Data Science

In this tutorial, we’ll be building a simple DCGAN in PyTorch and training it to generate handwritten digits. As part of this tutorial we’ll be discussing the PyTorch DataLoader and how to use it to feed real image data into a PyTorch neural network for training. PyTorch is the focus of this tutorial, so I’ll be assuming you’re familiar with how GANs work.
Python 3.7 or higher. Any lower and you’ll have to refactor the f-strings.PyTorch 1.5 Not sure how to install it? This might help.Matplotlib 3.1 or higherTwenty-two minutes or so of your time.
Python 3.7 or higher. Any lower and you’ll have to refactor the f-strings.
PyTorch 1.5 Not sure how to install it? This might help.
Matplotlib 3.1 or higher
Twenty-two minutes or so of your time.
It’s not required but I recommended reading my vanilla GAN tutorial first; it explains a couple things that this tutorial takes for granted. I also recommend that you do this tutorial on a computer with a CUDA GPU, or keep a hefty book of Sudokus at the ready.
Create a function G: Z → X where Z~N16(0, 1) and X~MNIST.
That is to say, train a GAN that takes 16-dimensional random noise and produces images that look like real samples from the MNIST dataset.
...let’s do a little housekeeping. If you haven’t already, install the required versions of Python and the above libraries. Then, create your project directory. I called mine DCGAN. In that directory, make a directory called data. Then, navigate to this GitHub repo and download mnist_png.tar.gz. This compressed file contains the MNIST dataset as 70000 individual png files. Of course, we could be using PyTorch’s built-in MNIST dataset, but then you wouldn’t learn how to actually load image data for training. Decompress the file and put the mnist_png directory into your data directory. Create a file called dcgan_mnist.py, and put it in your DCGAN directory. You project directories should look like this:
Finally, add the following to your dcgan_mnist.py script:
import osimport torchfrom torch import nnfrom torch import optimimport torchvision as tvfrom torchvision.datasets import ImageFolderfrom torch.utils.data import DataLoader
Alright, now we’re ready to start.
Add the following to your dcgan_mnist.py script:
The Generator inherits nn.Module, which is the base class for a PyTorch neural network. The Generator has three methods:
The constructor, which stores the instance variables and calls _init_layers. There’s not much to say here.
This method instantiates the PyTorch modules (or “layers”, as they’re called in other frameworks). These include:
A linear (“fully-connected”) module, for mapping the latent space to a 7×7×256=12544-dimensional space. As we will see in the forward method, this 12544-length tensor is reshaped to a (256, 7, 7) “image” tensor (channels×height×width). In PyTorch, unlike TensorFlow, channels come before the spatial dimensions.
A 1-dimensional batch normalization module, if specified.
A leaky ReLU module.
A 2-dimensional convolutional layer.
Two 2-dimensional transposed convolutional layers; these are used for upscaling the image. Notice how the out-channels of one convolutional layer are the in-channels of the next.
Two 2-dimensional batch normalization layers, if specified.
A Tanh module as the output activation. We’ll be rescaling our images to the range [-1, 1], so our generator output activation should reflect that.
These could have been instantiated in the __init__ method, but I like to keep the module instantiation separate from the constructor. It’s trivial for a model this simple, but helps keep the code simple as the model grows in complexity.
This is the method our Generator uses to generate samples from random noise. The input tensor is passed to the first module, the output of that is passed to the next module, the output of that is passed to the next module, etc. It’s fairly straightforward, but I’d like to draw your attention to two interesting features:
Note the line intermediate = intermediate.view((-1, 256, 7, 7)). Unlike Keras, PyTorch doesn’t use an explicit “reshape” module; instead, we manually reshape the Tensor using the PyTorch operation “view”. Other simple PyTorch operations can be applied during the forward pass as well, like multiplying a tensor by two, and PyTorch won’t bat an eye.Notice how there are if statements in the forward method. PyTorch uses a define-by-run strategy, which means that the computational graph is built on-the-fly during the forward pass. This makes PyTorch extremely flexible; there’s nothing stopping you from adding loops to your forward pass, or randomly selecting one of several modules to use.
Note the line intermediate = intermediate.view((-1, 256, 7, 7)). Unlike Keras, PyTorch doesn’t use an explicit “reshape” module; instead, we manually reshape the Tensor using the PyTorch operation “view”. Other simple PyTorch operations can be applied during the forward pass as well, like multiplying a tensor by two, and PyTorch won’t bat an eye.
Notice how there are if statements in the forward method. PyTorch uses a define-by-run strategy, which means that the computational graph is built on-the-fly during the forward pass. This makes PyTorch extremely flexible; there’s nothing stopping you from adding loops to your forward pass, or randomly selecting one of several modules to use.
Add the following to your dcgan_mnist.py script:
I won’t go into too much detail about this one since it’s very similar to the Generator but backwards. Take a good read through it and make sure you understand what it’s doing.
Add the following to your dcgan_mnist.py script:
Let’s step through the constructor line-by-line:
self.generator = Generator(latent_dim).to(device)self.discriminator = Discriminator().to(device)
The first two (non-docstring) lines of the constructor instantiate the Generator and the Discriminator, move them to the specified device, and store them as instance variables. The device is typically “cpu”, or “cuda” if you want to use a gpu.
self.noise_fn = noise_fn
Next, we store noise_fn as an instance variable; noise_fn is a function that takes an integer num as input and returns num latent vectors as a PyTorch tensor as output with shape (num, latent_dim). This PyTorch tensor must be on the specified device.
self.dataloader = dataloader
We store dataloader, a torch.utils.data.DataLoader object, as an instance variable; more on this later.
self.batch_size = batch_sizeself.device = device
Store the batch size and device as instance variables. Simple.
self.criterion = nn.BCELoss()self.optim_d = optim.Adam(self.discriminator.parameters(), lr=lr_d, betas=(0.5, 0.999))self.optim_g = optim.Adam(self.generator.parameters(), lr=lr_g, betas=(0.5, 0.999))
Set the loss function to binary cross-entropy, and instantiate Adam optimizers for the generator and the discriminator. PyTorch optimizers need to know what they’re optimizing. For the discriminator, this means all trainable parameters within the Discriminator network. Because our Discriminator class inherits from nn.Module, it has the parameters() method which returns all the trainable parameters in all its instance variables that are also PyTorch Modules. The same goes for the Generator.
self.target_ones = torch.ones((batch_size, 1), device=device)self.target_zeros = torch.zeros((batch_size, 1), device=device)
Targets for training, set to the specified device. Remember, the Discriminator is trying to classify real samples as 1 and generated samples as 0, while the Generator is trying to get the the Discriminator to misclassify generated samples as 1. We define and store them here so we don’t have to remake them with each training step.
A helper method for generating samples. Note that the no_grad context manager is used, which tells PyTorch not to keep track of gradients because this method is not used for training the network. Also note that, regardless of the specified device, the returned tensor is set to cpu, which is necessary for further use such as displaying samples or saving them to disk.
This method performs one training step of the generator and returns the loss as a float. Let’s step through it:
self.generator.zero_grad()
Clear the generator’s gradient. This is necessary, since PyTorch automatically keeps track of gradients and the computational network. We don’t want one training step influencing the next.
latent_vec = self.noise_fn(self.batch_size)generated = self.generator(latent_vec)classifications = self.discriminator(generated)loss = self.criterion(classifications, self.target_ones)
Get a batch of latent vectors, use them to generate samples, discriminate how realistic each samples is, then calculate the loss using the binary cross-entropy criterion. Note that by chaining these networks together we are creating a single computation graph starting from the latent vector, including the Generator and Discriminator networks, and ending at the loss.
loss.backward()self.optim_g.step()
One of the main benefits of PyTorch is that it automatically keeps track of the computational graph and its gradients. By calling the backward method on the loss, PyTorch applies backpropagation and calculates the gradients of the loss with respect to each parameter in the computation graph. By then calling the Generator’s optimizer’s step method, the Generator’s parameters (and only the Generator’s parameters) are nudged slightly in the negative direction of their gradient.
return loss.item()
Finally, we return the loss. It’s important to use the item method, so that we’re returning a float instead of a PyTorch tensor. If we had instead returned the tensor, the Python garbage collector wouldn’t be able to clean up the underlying computational graph, and we would quickly run out of memory.
This method is very similar to train_step_generator, but with two notable differences. First:
with torch.no_grad():    fake_samples = self.generator(latent_vec)
The context manager no_grad is used here to to tell PyTorch not to worry about keeping track of gradients. It’s not necessary, but cuts down on unnecessary computation. Second:
loss = (loss_real + loss_fake) / 2
This line is really cool. loss_real is the Discriminator’s loss for the real samples (and attached to it is its computational graph), and loss_fake is the loss (and graph) for the fake samples. PyTorch is able to combine these into one computational graph using the + operator. We then apply backpropagation and parameter updates to that combined computational graph. If you don’t think that’s awesomely simple, try rewriting this in another framework.
This function trains the generator and discriminator for one epoch, which is one pass over the entire dataset. We’ll come back to this after a brief detour.
Add the following code to your script:
This function builds, trains, and showcases the GAN.
import matplotlib.pyplot as pltfrom time import timebatch_size = 32epochs = 100latent_dim = 16
Import pyplot (for visualizing the generated digits) and time (for timing the training). Set the training batch size to 32, the number of epochs to 100, and the latent dimension to 16.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
This line checks if a cuda device is available. If it is, device is assigned that device; otherwise, device is assigned the cpu.
transform = tv.transforms.Compose([            tv.transforms.Grayscale(num_output_channels=1),            tv.transforms.ToTensor(),            tv.transforms.Normalize((0.5,), (0.5,))            ])
This composite transform is used by the dataloader to pre-process images. The MNIST dataset that we downloaded earlier are .png files; when PyTorch loads them from disk, they have to be processed so that our neural network can use them properly. The transforms are, in order:
Grayscale(num_output_channels=1): Convert the image to greyscale. When loaded, the MNIST digits are in RGB format with three channels. Greyscale reduces these three to one.
ToTensor(): Convert the image to a PyTorch Tensor, with dimensions channels × height × width. This also rescales the pixel values, from integers between 0 and 255 to floats between 0.0 and 1.0.
Normalize((0.5,), (0.5,)): Scale and translate the pixel values from the range [0.0, 1.0] to [-1.0, 1.0]. The first argument is μ and the second argument is σ, and the function applied to each pixel is:
The reason that μ and σ are one-tuples is that this transform is applied per-channel. The equivalent transformation to an RGB image would be Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
dataset = ImageFolder(            root=os.path.join("data", "mnist_png", "training"),            transform=transform            )
Here we create the dataset by specifying its root and the transformations to apply. This is used to create the DataLoader:
dataloader = DataLoader(dataset,            batch_size=batch_size,            shuffle=True,            num_workers=2            )
The DataLoader is an object that... well, it loads data from a dataset. Here we specify our batch size, tell the dataloader to shuffle the dataset between epochs, and use multiprocessing with two worker processes (if you’re using Windows and this causes issues, set num_workers to 0). You can iterate through this dataloader, and with each iteration it will return a tuple containing:
a PyTorch tensor of shape (32, 1, 28, 28) corresponding to a batch (32 samples) of greyscale (1 channel) MNIST images (28×28 pixels).a PyTorch tensor of shape (32,) of digits 0 through 9, corresponding to the label (digit) of that image. These class labels were taken from the directory structure, as all the zeros were in directory 0, all the ones in 1, etc.
a PyTorch tensor of shape (32, 1, 28, 28) corresponding to a batch (32 samples) of greyscale (1 channel) MNIST images (28×28 pixels).
a PyTorch tensor of shape (32,) of digits 0 through 9, corresponding to the label (digit) of that image. These class labels were taken from the directory structure, as all the zeros were in directory 0, all the ones in 1, etc.
noise_fn = lambda x: torch.rand((x, latent_dim), device=device)
A function for generating random, normally-distributed noise.
gan = DCGAN(latent_dim, noise_fn, dataloader, device=device)start = time()for i in range(10):    print(f"Epoch {i+1}; Elapsed time = {int(time() - start)}s")    gan.train_epoch()
Build and train the GAN.
Let’s revisit this one, now that we’ve discussed what a DataLoader is. The method is pretty self-explanatory, albeit verbose, but I’d like to focus on two lines:
for batch, (real_samples, _) in enumerate(self.dataloader):   real_samples = real_samples.to(self.device)
Here, we iterate through the dataloader. We wrap the dataloader in an enumerator so that we can keep track of the batch number, but as you can see the dataloader does indeed return a tuple as promised. We assign the batch of images tensor to real_samples, and ignore the labels since we don’t need them. Then, in the loop, we move real_samples to the specified device. It’s important that the input to the model and the model itself are on the same device; don’t worry if you forget to do this, PyTorch will certainly let you know! Also, don’t worry about the dataloader “running out”. Once we’ve iterated through the whole dataset, the loop will end, but if we try to iterate through it again it’ll start back at the beginning (shuffling the images first, because we specified that when we made the dataloader).
If you copy and pasted correctly, running the script should show you training stats for a couple of minutes followed by some generated digits. Hopefully, it looks something like this:
If they look terrible and your loss exploded, try running it again (GANs are notoriously unstable). If it still doesn’t work, drop a comment down below and we’ll see if we can’t debug it.
Just for fun, I modified the script to see what the Generator was capable of after every 10 training steps. Here were the results.
I think that’s pretty good for only 1000 steps. Here’s the loss over those training steps, split into 10 step “epochs”.
The DCGAN described in this tutorial is obviously very simple, but it should be enough to get you started implementing more complex GANs in PyTorch.
Can you modify this script to make a Conditional GAN before I make a tutorial about it?
The full script is available here.
All images not cited are my own. Feel free to use them, but please cite this article ❤