Write your first Generative Adversarial Network Model on PyTorch
By: Mikhail Raevskiy
Tags: Data Science, Deep Learning, Machine Learning, Python, Programming

Generative adversarial networks (abbreviated GAN) are neural networks that can generate images, music, speech, and texts similar to those that humans do. GANs have become an active research topic in recent years. Facebook AI Lab Director Yang Lekun called adversarial learning “the most exciting machine learning idea in the last 10 years.” Below we will explore how GANs work and create two models using the PyTorch deep learning framework.

What is a Generative Adversarial Network?

A generative adversarial network (GAN) is a machine learning model that can simulate a given data distribution. The model was first proposed in a 2014 NeurIPS paper by deep learning expert Ian Goodfellow and colleagues.

GAN learning process

GANs consist of two neural networks, one of which is trained to generate data, and the other is trained to distinguish simulated data from real data (hence the “adversarial” nature of the model). Generative adversarial networks show impressive results in terms of image and video generation:

transfer of styles (CycleGAN) — the transformation of one image in accordance with the style of other images (for example, paintings by a famous artist);

Human Face Generation (StyleGAN), realistic examples are available at This Person Does Not Exist.

GANs and other data-generating structures are called generative models as opposed to more widely studied discriminative models. Before diving into GANs, let’s look at the differences between the two types of models.

Comparison of Discriminative and Generative Machine Learning models

Discriminative models are used for most supervised learning problems for classification or regression. As an example of a classification problem, suppose you want to train a handwritten digit image recognition model. To do this, we can use a labeled dataset containing photographs of handwritten numbers to which the numbers themselves are correlated.

Training is reduced to setting the parameters of the model using a special algorithm that minimizes the loss function. The loss function is a criterion for the discrepancy between the true value of the estimated parameter and its expectation. After the learning phase, we can use the model to classify a new (previously not considered) handwritten digit image by matching the most likely digit to the input image.

Discriminative model training scheme

The discriminative model uses training data to find the boundaries between classes. The found boundaries are used to distinguish new inputs and predict their class. Mathematically, discriminative models study the conditional probability P (y | x) of an observation y for a given input x.

Discriminative models are not only neural networks but also logistic regression and support vector machine (SVM).

While discriminative models are used for supervised learning, generative models typically use a RAW data set, that is, may be seen as a form of unsupervised learning. So, using a dataset of handwritten numbers, you can train a generative model to generate new images.

In contrast to discriminative models, generative models study the properties of the probability function P (x) of the input data 𝑥 . As a result, they do not generate a prediction, but a new object with properties akin to the training dataset.

Besides GAN, there are other generative architectures:

GANs have gained a lot of attention recently for their impressive results in visual content generation. Let’s dwell on the device of generative adversarial networks in more detail.

The architecture of generative adversarial neural networks

A generative adversarial network, as we have already understood, is not one network, but two: a generator and a discriminator. The role of the generator is to generate a dataset based on a real sample that resembles real data. The discriminator is trained to estimate the probability that the sample is obtained from real data and not provided by a generator. Two neural networks play cat and mouse: the generator tries to trick the discriminator, and the discriminator tries to better identify the generated samples.

To understand how GAN training works, consider a toy example with a dataset consisting of two-dimensional samples (x1, x2), with x1 ranging from 0 to 2π and x2=sin(x1).

Dependence x2 on x1

The general structure of the GAN for generating pairs (x̃1, x̃2) resembling points from a dataset is shown in the following figure.

General GAN ​​structure

The generator receives as input pairs of random numbers (z1, z2), transforming them so that they resemble examples from a real sample. The structure of the neural network can be any, for example, a multilayer perceptron or a convolutional neural network. G G

The discriminator inputs alternately samples from the training dataset and simulated samples provided by the generator. The role of the discriminator is to assess the likelihood that the input data belongs to a real dataset. That is, training is performed in such a way that it gives out, receiving a real sample, and for the generated sample. D G D 1 0

As in the case with the generator, you can choose any structure of the neural network, taking into account the sizes of the input and output data. In this example, the input is 2D and the output is a scalar ranging from 0 to 1. D

Mathematically, the GAN learning process consists of a minimax game of two players, in which it is adapted to minimize the error of the difference between the real and the generated sample, and adapted to maximize the probability of making an error. D G D

At each stage of training, the parameters of the models and are updated. To train, at each iteration, we mark a sample of real samples with ones and a sample of generated samples created with zeros. Thus, a normal supervised learning approach can be used to update the parameters as shown in the diagram. D G D G D

Discriminator training process

For each batch of training data containing tagged real and generated samples, we update the set of model parameters D , minimizing the loss function. After the parameters are D updated, we train to G generate better samples. The set of parameters is D "frozen" during the training of the generator.

Generator learning process

When it starts to generate samples so well that it is “fooled”, the output probability tends to one — it considers that all samples belong to the original sample. G D D

Now that we know how GAN works, we are ready to implement our own neural network using PyTorch.

Your first generative adversarial network

As the first experiment with generative adversarial networks, we will implement the above example with a harmonic function. To work with the example, we will use the popular PyTorch library, which can be installed using the instructions. If you’re seriously interested in Data Science, you may have already used the Anaconda distribution and the conda package and environment management system . Note that the environment makes the installation process easier.

Installing PyTorch with, first create an environment and activate it: conda

$ conda create --name gan

$ conda activate gan

This creates an environment conda named gan . Inside the created environment, you can install the necessary packages:

$ conda install -c pytorch pytorch=1.4.0

$ conda install matplotlib jupyter

Since PyTorch is an actively developing environment, the API may change in new versions. Code examples have been verified for version 1.4.0.

We will use matplotlib to work with graphs.

When using Jupyter Notebook, you need to register the environment so that you can create notebooks using this environment as a kernel. To do this, in the activated environment, run the following command: conda gan gan

$ python -m ipykernel install --user --name gan

Let’s start by importing the required libraries:

Here we are importing the PyTorch ( torc ) library. We import the component separately from the library for more compact handling. The built-in library is only needed to get the value of the constant, and the tool mentioned above is for building dependencies. nn math pi matplotlib

It is good practice to temporarily secure the random number generator so that the experiment can be replicated on another machine. To do this in PyTorch, run the following code:

We use the 111 number to initialize the random number generator. We will need a generator to set the initial weights of the neural network. Despite the random nature of the experiment, its course will be reproducible.

Preparing data for GAN training

The training set consists of pairs of numbers (x1, x2) — such that x2 corresponds to the sine value of x1 for x1 in the range from 0 to 2π. Training data can be obtained as follows:

Here we compile a training dataset of 1024 pairs (x1, x2). Then we initialize with zeros — a matrix of 1024 rows and 2 columns. train_data

The first column is filled with random values ​​in the range from 0 to 2π. We calculate the values ​​of the second column as the sine of the first. train_data

We then formally need an array of labels, which we pass to the PyTorch data loader. Since the GAN implements unsupervised learning, the labels can be anything. train_labels

Finally, we create a list of tuples from and. train_data train_labels train_set

Let’s display the data for training by plotting each point (x1, x2):

Construction result

Let’s create a data loader named train_loader that will shuffle data from train_set , returning packets of 32 samples ( batch_size ) used to train the neural network:

The data is ready, now you need to create the discriminator and GAN neural networks.

GAN Discriminator Implementation

In PyTorch, neural network models are represented by classes that inherit from a class. If you are new to OOP, the article “An Introduction to Object-Oriented Programming (OOP) in Python” will suffice to understand what is happening. nn.Module

The discriminator is a two-dimensional input and one-dimensional output model. It takes a sample from real data or from a generator and provides the probability that the sample is from real training data. The code below shows how to create a discriminator class.

A standard class method is used to build a neural network model . Inside this method, we first call to run the corresponding method of the inherited class . A multilayer perceptron is used as the architecture of the neural network . Its structure is specified in layers using . The model has the following characteristics: __init__() super().__init__() __init__() nn.Module nn.Sequential()

two-dimensional entrance;

the first hidden layer consists of 256 neurons and has a ReLU activation function ;

in the subsequent layers, the number of neurons decreases to 128 and 64. The output has a sigmoidal activation function, which is characteristic of representing the probability ( Sigmoid );

); to avoid overfitting, after the first, second and third hidden layers, a part of the neurons is dropped ( Dropout ).

For the convenience of inference, a method is also created in the class. Here corresponds to the input of the model. In this implementation, the output is obtained by feeding the input into the model we have defined without preprocessing. forward() x x

After declaring the discriminator class, create an instance of it:

discriminator = Discriminator()

GAN generator implementation

In generative adversarial networks, a generator is a model that takes as input some sample from a space of hidden variables that resemble the data in the training set. In our case, this is a 2D input model that will receive random points (z1, z2), and a 2D output that produces points (x̃1, x̃2) that look like the points from the training data.

The implementation is similar to what we wrote for the discriminator. First, you need to create a class that inherits from, then define the architecture of the neural network, and finally create an instance of the object : Generator nn.Module Generator

The generator includes two hidden layers with 16 and 32 neurons with the ReLU activation function, and at the output a layer with two neurons with a linear activation function. Thus, the output will consist of two elements ranging from −∞ to + ∞ , which will represent (x̃1 , x̃2). That is, initially we do not impose any restrictions on the generator — it must “learn everything by itself.”

Now that we have defined the models for the discriminator and generator, we are ready to start training.

Train GAN Models

Before training the models, you need to configure the parameters that will be used in the training process:

What’s going on here:

We set the learning rate, which we will use to adapt the network weights. lr We set the number of epochs, which determines how many repetitions of the training process will be performed using the entire dataset. num_epochs To the variable, we assign the function of the logistic loss function (binary cross-entropy). This is the loss function that we will use to train the models. It is suitable both for training a discriminator (its task is reduced to a binary classification) and for a generator since it feeds its output to the input of the discriminator. loss_function BCELoss()

The rules for updating weights (training the model) in PyTorch are implemented in a module. We will use Adam’s stochastic gradient descent algorithm to train discriminator and generator models. To create optimizers, run the following code: torch.optim torch.optim

Finally, it is necessary to implement a training cycle in which samples of the training sample are fed to the model input, and their weights are updated, minimizing the loss function:

Here, at each training iteration, we update the discriminator and generator parameters. As is usually done for neural networks, the training process consists of two nested loops: the outer one for the training epochs, and the inner one for the packets within each epoch. In the inner loop, it all starts with preparing data for training the discriminator:

We get real samples of the current batch from the data loader and assign them to a variable. Note that the first dimension in the array dimension has the number of elements equal to. This is the standard way of organizing data in PyTorch, where each tensor row represents one sample from the package. real_samples batch_size

Use to create labels with a value of 1 for real samples and assign labels to a variable. torch.ones() real_samples_labels

We generate samples by storing random data, which we then pass to the generator for receiving. We use zeros for the labels of the generated samples, which we store in latent_space_samples generate_samples torch.zeros() generate_samples_labels

It remains to combine the real and generated samples and labels and save respectively in and all_samples all_samples_labels

In the next block, we train the discriminator:

In PyTorch, it is important to clear the gradient values ​​at every step of the training. We do this using the method zero_grad()

We calculate the output of the discriminator using the training data all_samples

Calculate the value of the loss function using the output at and labels output_discriminator all_samples_labels

Calculate the gradients to update the weights with loss_discriminator.backward()

Find the updated discriminator weights by calling optimizer_discriminator.step()

We prepare the data for training the generator. We use two columns to match the 2D data at the generator input. latent_space_samples batch_size

We train the generator:

We clean up the gradients using the method. zero_grad()

We pass it on to the generator and save its output to latent_space_samples generate_samples

We pass the generator output to the discriminator and save its output, which will be used as the output of the entire model. output_discriminator_generated

Calculate the loss function using the output of the classification system stored in and labels equal to 1. output_discriminator_generated real_samples_labels

Calculating gradients and updating generator weights. Remember that when we train the generator, we are keeping the discriminator weights frozen.

Finally, in the last lines of the loop, the values ​​of the discriminator and generator loss functions are output at the end of every tenth epoch.

Checking samples generated by GAN

Generative adversarial networks are designed to generate data. Thus, after the training process is complete, we can call the generator to get new data:

Let’s plot the generated data and check how similar it is to the training data. Before plotting a graph for the generated samples, you need to apply a method detach() to get the necessary data from the PyTorch computational graph:

Results of building the generated dataset

The distribution of the generated data is very similar to real data — the original sine. The animation of the evolution of learning can be viewed here .

At the beginning of the training process, the distribution of the generated data is very different from the real data. But as it learns, the generator learns the real data distribution, as if adjusting to it.

Now that we have implemented the first model of a generative adversarial network, we can move on to a more practical example of generating images.

Handwritten Digit Generator with GAN

In the following example, we will use GAN to generate images of handwritten numbers. To do this, we will train the models using the MNIST dataset of handwritten numbers. This standard dataset is included in the package torchvision

First, in the activated environment, you need to install : gan torchvision

$ conda install -c pytorch torchvision=0.5.0

Again, here we are specifying the specific version just like we did with PyTorch to ensure that the code examples run. torchvision

We start by importing the required libraries:

In addition to the libraries that we imported earlier, we will also need to transform the information stored in image files. torchvision torchvision.transforms

Since the training set includes images in this example, the models will be more complex and the training will take significantly longer. When training in a central processing unit (CPU), one epoch will take about two minutes. It will take about 50 epochs to get an acceptable result, so the total training time using the processor is about 100 minutes.

A graphics processing unit (GPU) can be used to reduce training time.

To make the code work regardless of the characteristics of the computer, let’s create an object that will point either to the central processor or (if available) to the graphics processor: device

The environment is configured, let’s prepare a dataset for training.

Preparing the MNIST dataset

The MNIST dataset consists of images of handwritten digits 0 through 9. The images are in grayscale and are 28 × 28 pixels in size. To use them with PyTorch, you need to do some transformations. To do this, we define the function used when loading data: transform

The function has two parts:

transforms.ToTensor() converts the data into a PyTorch tensor. transforms.Normalize() converts a range of tensor coefficients.

The original coefficients are given by the function range from 0 to 1. Since the images have a black background, most of the coefficients are 0. transforms.ToTensor()

The function changes the range of coefficients to transforms.Normalize() [ - 1 , 1 ][−1,1] , subtracting 0.5 from the original odds and dividing the result by 0.5. The transformation reduces the number of elements in the input samples to zero. This helps in training the models.

We can now load the training data by calling : torchvision.datasets.MNIST

The argument ensures that the first time you run the code, the MNIST dataset will be loaded and saved in the current directory as specified in the argument. download = True root

We created so that we can create a data loader as we did before: train_set

Let’s use matplotlib for the selective plotting of data. Well suited as a palette cmap = gray_r . The numbers will be displayed in black on a white background:

As you can see, the dataset contains numbers with different handwriting. As the GAN learns the distribution of the data, it also generates numbers with different handwriting styles.

We have prepared training data, we can implement discriminator and generator models.

Discriminator and generator implementation

In this case, the discriminator is a multilayer perceptron neural network, which takes an image of 28 × 28 pixels and finds the probability that the image belongs to real training data.

To introduce the image coefficients into the perceptron neural network, it is necessary to vectorize them so that the neural network receives a vector consisting of 784 coefficients (28 × 28 = 784).

Vectorization occurs in the first line of the method — the call transforms the form of the input tensor. Initial tensor form forward() x.view() 𝑥, where 32 is the batch size. After transformation, the form 32 × 1 × 28 × 28 𝑥x becomes equal, with each row representing the image coefficients of the training set. 32 × 784

To run a discriminator model using a GPU, you need to instantiate it and associate it with a device object using the method : to()

discriminator = Discriminator().to(device=device)

The generator will create more complex data than the previous example. Therefore, it is necessary to increase the size of the input data used for initialization. Here we are using a 100-dimensional input and output with 784 coefficients. The result is organized as a 28x28 tensor representing the image.

The output coefficients must be in the range from -1 to 1. Therefore, at the output of the generator, we use a hyperbolic activation function. In the last line, we instantiate the generator and associate it with the device object. Tanh()

It remains only to train the model.

Model training

To train models, you need to define training parameters and optimizers:

We are reducing the learning rate compared to the previous example. To shorten the training time, set the number of epochs to 50.

The learning loop is similar to the one we used in the previous example:

Checking Generated GAN Samples

Let’s generate some samples of “handwritten numbers”. To do this, we pass the generator an initiating set of random numbers:

To build the generated samples, you need to move the data back to the central processor, if it was processed on the GPU. To do this, just call the method cpu() . As before, before plotting the data, you need to call the method detach() :

The output should be numbers that resemble training data.

Result of generating images

After fifty epochs of learning, there are several numbers, as if written by a human hand. Results can be improved with longer training times (with more epochs). As in the previous example, you can visualize the evolution of training by using a fixed tensor of the input and feeding it to the generator at the end of each epoch (animation of the evolution of training).

At the beginning of the training process, the generated images are completely random. As it learns, the generator learns the distribution of real data, and after about twenty epochs some of the generated digit images already resemble real data.

Conclusion

Congratulations! You have learned how to implement your own generative adversarial network. We first built a toy example to understand the structure of the GAN, and then looked at a network for generating images from the available sample data.

Despite the complexity of the GAN topic, machine learning frameworks like PyTorch make implementation very easy.