Deep Convolutional Generative Adversarial Network (DCGAN)

September 26, 2020

---

Generative Adversarial Networks GANs are a type of generative models. The objective of generative models is to learn the training data distribution, so as to enable generation of new data, through sampling from the same distribution. The other popular types of generative models are Variational Auto Encoders VAEs and Normalizing Flow based models e.g. NICE, Glow. Deep Convolutional Generative Adversarial Network DCGAN, as the name suggests is a GAN. The distinguishing feature is that it uses convolutions in the discriminator and transposed convolutions in the generator.

GANs have a very unique training procedure. In the first phase, we train the discriminator and in the second phase, we train the generator. While training DCGAN on MNIST, a single training step (forward and backward propagation) invokes about 500 GPU CUDA kernels. The best way to understand the GAN training procedure and GPU performance is to understand every CUDA kernel i.e. which layer of the network invoked the kernel, with what arguments (tensor shapes and datatypes) and in which direction (forward propagation or backward propagation).

In this post, we will categorize every kernel used in the training of DCGAN. All the information in the tables below was obtaining using Nvidia's PyTorch Profiler, PyProf, on a Turing T4 GPU. The information below is only a subset of what is provided by PyProf. The code and instructions for obtaining a detailed profile are here. Note that different GPUs will have slightly different kernel names e.g. volta_* as opposed to turing_*.

The code for DCGAN was obtained from PyTorch DCGAN Tutorial and modified to use MNIST. The image below shows the output of the generator on a fixed noise at the beginning and at the end of epochs 1 through 10.

Model Parameters

The parameters in the profiled code are as follows.

# Size of training images.
image_size = 64

# Channels in the training images.
# 3 for color images, 1 for MNIST.
nc = 1

# Size of latent vector (i.e. size of generator input).
nz = 100

# Size of feature maps in generator.
ngf = 64

# Size of feature maps in discriminator.
ndf = 64

batch_size = 128

GPU Kernels

The tables below show the GPU kernels invoked in 1 training step. For every GPU kernel we show the direction (fprop, bprop), name of the layer, name of the operation, and the input tensor shapes / matrix dimensions for the operation. PyProf provides a lot of additional information for every GPU kernel e.g. grid dimensions, block dimensions, silicon time, datatypes, flops, bytes, tensor core usage and so on.

GAN training consists of two parts. Part 1, where we train the discriminator and part 2, where we train the generator.

Part 1: Train the Discriminator

Zero out the discriminator gradients.

At the beginning of part 1, we zero out the gradients of the discriminator.

Discriminator: Forward propagation on real images from the dataset.

We pick batch_size i.e. 128 images from the MNIST dataset and pass them through the discriminator. The target label for these images is set to 1 (kernel 12). The discriminator consists of 5 convolution layers. Kernels 13-15 correspond to the first convolution and activation. The input shape (N,C,H,W) is (128,1,64,64), the number of filters K is 64, the kernel size (R,S) is (4,4), the padding (ph,pw) is (1,1) and the horizontal and vertical stride (U,V) is (2,2). The output shape (N,K,P,Q) is (128,64,32,32). Kernels 16-20 correspond to the second convolution, batch norm and activation. The output shape is (128,128,16,16). Kernels 21-25 correspond to the third convolution, batch norm and activation. The output shape is (128,256,8,8). Kernels 26-30 correspond to the fourth convolution, batch norm and activation. The output shape is (128,512,4,4). Kernels 31-33 correspond to the fifth convolution and activation. Note that K=1, therefore, the output shape is (128,1,1,1). Kernels 34-35 calculate the loss with respect to misclassification of real images from the dataset.

Discriminator: Backward propagation (real images).

We now perform back propagation through the discriminator and calculate the gradients. Kernels 37-38 correspond to bprop through the loss layer. Kernels 39-46 correspond to bprop (data gradient and weight gradient) through the fifth convolution layer. Kernels 48,49,52-57 correspond to bprop through the fourth convolution layer. Kernels 59,60,63-70 correspond to bprop through the third convolution layer. Kernels 72,73,76-78 correspond to bprop through the second convolution layer. Kernels 80,81 correspond to bprop through the first convolution layer. Note that the first convolution layer requires only a weight gradient and not a data gradient. Kernels with the op add_, most likely correspond to gradient accumulation i.e. adding the gradients to the previously zeroed out gradient tensors. Kernel 83 calculates the average loss (for reporting).

Generator.

We now create batch_size i.e. 128 fake images using the generator. The generator consists of 5 transposed convolutions which progressively increase the image size from [nz,1,1] to [nc,ngf,ngf] i.e. from [100,1,1] to [1,64,64]. Kernel 84 creates a random tensor of shape [batch_size, nz]. Kernels 85-90 correspond to the first transposed convolution, batch norm and activation. The output shape is [128,512,4,4]. Kernels 91-95 correspond to the second transposed convolution, batch norm and activation. The output shape is [128,256,8,8]. Kernels 96-102 correspond to the third transposed convolution, batch norm and activation. The output shape is [128,128,16,16]. Kernels 103-107 correspond to the fourth transposed convolution, batch norm and activation. The output shape is [128,64,32,32]. Kernels 108-110 correspond to the fifth transposed convolution and activation. The output shape is [128,1,64,64].

Discriminator: Forward propagation on fake images from the generator.

The fake images from the generator are now fed to the discriminator. The target label for these images is set to 0 (kernel 111). Kernels 112 through 134 are the same as when the discriminator was fed real images from the dataset.

Discriminator: Backward propagation (fake images).

We now perform back propagation through the discriminator again and calculate and accumulate the gradients. Kernels 135 through 182 are the same as kernels 36 through 83. Kernel 183 adds the losses from real and fake images (for reporting).

Discriminator Optimizer

After calculating and summing up the gradients from the real and fake images, we apply the Adam optimizer on the discriminator weights (parameters). The discriminator has 11 parameters, 1 for each of the 5 convolutions and 2 for each of 3 batch norms (see kernels 1-11). Each call to the Adam optimizer invokes 8 kernels, for a total of 88 kernels (184 through 271). This is not an optimized implementation and one can use the fused Adam implementation from Nvidia Apex.

Part 2: Train the Generator

Zero out the generator gradients.

At the beginning of part 2, we zero out the gradients of the generator.

Discriminator: Forward propagation on fake images from the generator.

In part 1, for training the discriminator, we generated batch_size i.e. 128 fake images using the generator. Now for training the generator, we reuse those images and feed them through the discriminator. However, this time the target label for these images is set to 1 (kernel 285). Kernels 286 through 308 are the same as before.

Discriminator: Backward propagation (fake images).

We now perform back propagation through the discriminator. Kernels 309 through 357 are the same as kernels 36 through 82. The only difference is that we now calculate the data gradient for the first convolution layer as well, which results in 2 additional kernels. Ideally, in part 2, since we don't update the discrimator parameters, we only need the data gradients and not the weight gradients.

Generator: Backward propagation.

We now perform back propagation through the generator. Kernels 358-361 correspond to bprop through the fifth transposed convolution layer. Kernels 363,364,367-369 correspond to bprop through the fourth transposed convolution layer. Kernels 371,372,375-380 correspond to bprop through the third transposed convolution layer. Kernels 382,383,386-391 correspond to bprop through the second transposed convolution layer. Kernels 393,394,397-400 correspond to bprop through the first transposed convolution layer. Kernels with the op add_, most likely correspond to gradient accumulation i.e. adding the gradients to the previously zeroed out gradient tensors. Kernel 402 calculates the average loss (for reporting).

Generator Optimizer

The last step is to apply the Adam optimizer on the generator weights. The generator has 13 parameters, 1 for each of the 5 transposed convolutions and 2 for each of the 4 batch norms (see kernels 272-284). Each call to the Adam optimizer invokes 8 kernels, for a total of 104 kernels (403 through 506). This is not an optimized implementation and one can use the fused Adam implementation from Nvidia Apex.