Introduction to PyTorch

This notebook is a part of AI for Beginners Curricula. Visit the repository for complete set of learning materials.

Neural Frameworks

We have learnt that to train neural networks you need:

• Quickly multiply matrices (tensors)
• Compute gradients to perform gradient descent optimization

What neural network frameworks allow you to do:

• Operate with tensors on whatever compute is available, CPU or GPU, or even TPU
• Automatically compute gradients (they are explicitly programmed for all built-in tensor functions)

Optionally:

• Neural Network constructor / higher level API (describe network as a sequence of layers)
• Simple training functions (fit, as in Scikit Learn)
• A number of optimization algorithms in addition to gradient descent
• Data handling abstractions (that will ideally work on GPU, too)

Most Popular Frameworks

• Tensorflow 1.x - first widely available framework (Google). Allowed to define static computation graph, push it to GPU, and explicitly evaluate it
• PyTorch - a framework from Facebook that is growing in popularity
• Keras - higher level API on top of Tensorflow/PyTorch to unify and simplify using neural networks (Francois Chollet)
• Tensorflow 2.x + Keras - new version of Tensorflow with integrated Keras functionality, which supports dynamic computation graph, allowing to perform tensor operations very similar to numpy (and PyTorch)

In this Notebook, we will learn to use PyTorch. You need to make sure that you have recent version of PyTorch installed - to do it, follow the instructions on their site. It is normally as simple as doing

pip install torch torchvision

or

conda install pytorch -c pytorch

Basic Concepts: Tensor

Tensor is a multi-dimensional array. It is very convenient to use tensors to represent different types of data:

• 400x400 - black-and-white picture
• 400x400x3 - color picture
• 16x400x400x3 - minibatch of 16 color pictures
• 25x400x400x3 - one second of 25-fps video
• 8x25x400x400x3 - minibatch of 8 1-second videos

Simple Tensors

You can easily create simple tensors from lists of np-arrays, or generate random ones:

You can use arithmetic operations on tensors, which are performed element-wise, as in numpy. Tensors are automatically expanded to required dimension, if needed. To extract numpy-array from tensor, use .numpy():

In-place and out-of-place Operations

Tensor operations such as +/add return new tensors. However, sometimes you need to modify the existing tensor in-place. Most of the operations have their in-place counterparts, which end with _:

This is how we can compute the sum or all rows in a matrix in a naive way:

But it is much better to use

You can read more on PyTorch tensors in the official documentation

Computing Gradients

For back propagation, you need to compute gradients. We can set any PyTorch Tensor's attribute requires_grad to True, which will result in all operations with this tensor being tracked for gradient calculations. To compute the gradients, you need to call backward() method, after which the gradient will become available using grad attribute:

To be more precise, PyTorch automatically accumulates gradients. If you specify retain_graph=True when calling backward, computational graph will be preserved, and new gradient is added to the grad field. In order to restart computing gradients from scratch, we need to reset grad field to 0 explicitly by calling zero_():

To compute gradients, PyTorch creates and maintains compute graph. For each tensor that has the requires_grad flag set to True, PyTorch maintains a special function called grad_fn, which computes the derivative of the expression according to chain differentiation rule:

Here c is computed using mean function, thus grad_fn point to a function called MeanBackward.

In most of the cases, we want PyTorch to compute gradient of a scalar function (such as loss function). However, if we want to compute the gradient of a tensor with respect to another tensor, PyTorch allows us to compute the product of a Jacobian matrix and a given vector.

Suppose we have a vector function $\vec{y}=f(\vec{x})$, where $\vec{x}=\langle x_1,\dots,x_n\rangle$ and $\vec{y}=\langle y_1,\dots,y_m\rangle$, then a gradient of $\vec{y}$ with respect to $\vec{x}$ is defined by a Jacobian:

Instead of giving us access to the whole Jacobian, PyTorch computes the product $v^T\cdot J$ of Jacobian with some vector $v=(v_1 \dots v_m)$. In order to do that, we need to call backward and pass v as an argument. The size of v should be the same as the size of the original tensor, with respect to which we compute the gradient.

More on computing Jacobians in PyTorch can be found in official documentation

Example 0: Optimization Using Gradient Descent

Let's try to use automatic differentiation to find a minimum of a simple two-variable function $f(x_1,x_2)=(x_1-3)^2+(x_2+2)^2$. Let tensor x hold the current coordinates of a point. We start with some starting point $x^{(0)}=(0,0)$, and compute the next point in a sequence using gradient descent formula:

Here $\eta$ is so-called learning rage (we will denote it by lr in the code), and $\nabla f = (\frac{\partial f}{\partial x_1},\frac{\partial f}{\partial x_2})$ - gradient of $f$.

To begin, let's define starting value of x and the function f:

Now let's do 15 iterations of gradient descent. In each iteration, we will update x coordinates and print them, to make sure that we are approaching the minimum point at (3,-2):

Example 1: Linear Regression

Now we know enough to solve the classical problem of Linear regression. Let's generate small synthetic dataset:

Linear regression is defined by a straight line $f_{W,b}(x) = Wx+b$, where $W, b$ are model parameters that we need to find. An error on our dataset $\{x_i,y_u\}_{i=1}^N$ (also called loss function) can be defined as mean square error:

Let's define our model and loss function:

We will train the model on a series of minibatches. We will use gradient descent, adjusting model parameters using the following formulae:

Let's do the training. We will do several passes through the dataset (so-called epochs), divide it into minibatches and call the function defined above:

We now have obtained optimized parameters $W$ and $b$. Note that their values are similar to the original values used when generating the dataset ($W=2, b=1$)

Computations on GPU

To use GPU for computations, PyTorch supports moving tensors to GPU and building computational graph for GPU. Traditionally, in the beginning of our code we define available computation device device (which is either cpu or cuda), and then move all tensors to this device using a call .to(device). We can also create tensors on the specified device upfront, by passing the parameter device=... to tensor creation code. Such code works without changes both on CPU and GPU:

Example 2: Classification

Now we will consider binary classification problem. A good example of such a problem would be a tumour classification between malignant and benign based on it's size and age.

The core model is similar to regression, but we need to use different loss function. Let's start by generating sample data:

Training One-Layer Perceptron

Let's use PyTorch gradient computing machinery to train one-layer perceptron.

Our neural network will have 2 inputs and 1 output. The weight matrix $W$ will have size $2\times1$, and bias vector $b$ -- $1$.

To make our code more structured, let's group all parameters into a single class:

Note that we use W.data.zero_() instead of W.zero_(). We need to do this, because we cannot directly modify a tensor that is being tracked using Autograd mechanism.

Core model will be the same as in previous example, but loss function will be a logistic loss. To apply logistic loss, we need to get the value of probability as the output of our network, i.e. we need to bring the output $z$ to the range [0,1] using sigmoid activation function: $p=\sigma(z)$.

If we get the probability $p_i$ for the i-th input value corresponding to the actual class $y_i\in\{0,1\}$, we compute the loss as $\mathcal{L_i}=-(y_i\log p_i + (1-y_i)log(1-p_i))$.

In PyTorch, both those steps (applying sigmoid and then logistic loss) can be done using one call to binary_cross_entropy_with_logits function. Since we are training our network in minibatches, we need to average out the loss across all elements of a minibatch - and that is also done automatically by binary_cross_entropy_with_logits function:

The call to binary_crossentropy_with_logits is equivalent to a call to sigmoid, followed by a call to binary_crossentropy

To loop through our data, we will use built-in PyTorch mechanism for managing datasets. It is based on two concepts:

• Dataset is the main source of data, it can be either Iterable or Map-style
• Dataloader is responsible for loading the data from a dataset and splitting it into minibatches.

In our case, we will define a dataset based on a tensor, and split it into minibatches of 16 elements. Each minibatch contains two tensors, input data (size=16x2) and labels (a vector of length 16 of integer type - class number).

Now we can loop through the whole dataset to train our network for 15 epochs:

Obtained parameters:

To make sure our training worked, let's plot the line that separates two classes. Separation line is defined by the equation $W\times x + b = 0.5$

Not let's compute the accuracy on the validation dataset:

Let's explain what is going on here:

• pred is the vector of predicted probabilities for the whole validation dataset. We compute it by running original validation data valid_x through our network, and applying sigmoid to get probabilities.
• pred.view(-1) creates a flattened view of the original tensor. view is similar to reshape function in numpy.
• pred.view(-1)>0.5 returns a boolean tensor or truth value showing the predicted class (False = class 0, True = class 1)
• Similarly, torch.tensor(valid_labels)>0.5) creates the boolean tensor of truth values for validation labels
• We compare those two tensors element-wise, and get another boolean tensor, where True corresponds to correct prediction, and False - to incorrect.
• We convert that tensor to floating point, and take it's mean value using torch.mean - that is the desired accuracy

Neural Networks and Optimizers

In PyTorch, a special module torch.nn.Module is defined to represent a neural network. There are two methods to define your own neural network:

• Sequential, where you just specify a list of layers that comprise your network
• As a class inherited from torch.nn.Module

First method allows you to specify standard networks with sequential composition of layers, while the second one is more flexible, and gives an opportunity to express networks of arbitrary complex architectures.

Inside modules, you can use standard layers, such as:

• Linear - dense linear layer, equivalent to one-layered perceptron. It has the same architecture as we have defined above for our network
• Softmax, Sigmoid, ReLU - layers that correspond to activation functions
• There are also other layers for special network types - convolution, recurrent, etc. We will revisit many of them later in the course.

Most of the activation function and loss functions in PyTorch are available in two form: as a function (inside torch.nn.functional namespace) and as a layer (inside torch.nn namespace). For activation functions, it is often easier to use functional elements from torch.nn.functional, without creating separate layer object.

If we want to train one-layer perceptron, we can just use one built-in Linear layer:

As you can see, parameters() method returns all the parameters that need to be adjusted during training. They correspond to weight matrix $W$ and bias $b$. You may note that they have requires_grad set to True, because we need to compute gradients with respect to parameters.

PyTorch also contains built-in optimizers, which implement optimization methods such as gradient descent. Here is how we can define a stochastic gradient descent optimizer:

Using the optimizer, our training loop will look like this:

You may notice that to apply our network to input data we can use net(x) instead of net.forward(x), because nn.Module implements Python __call__() function

Taking this into account, we can define generic train function:

Defining Network as a Sequence of Layers

Now let's train multi-layered perceptron. It can be defined just by specifying a sequence of layers. The resulting object will automatically inherit from Module, e.g. it will also have parameters method that will return all parameters of the whole network.

We can train this multi-layered network using the function train that we have defined above:

Defining a Network as a Class

Using a class inherited from torch.nn.Module is a more flexible method, because we can define any computations inside it. Module automates a lot of things, eg. it automatically understands all internal variables that are PyTorch layers, and gathers their parameters for optimization. You just need to define all layers of the network as members of the class:

Task 1: Plot the graphs of loss function and accuracy on training and validation data during training

Task 2: Try to solve MNIST classificiation problem using this code. Hint: use crossentropy_with_logits as a loss function.

Defining a Network as PyTorch Lightning Module

Let's wrap the written PyTorch model code in PyTorch Lightining module. This allows to work with your model more conveniently and flexibly using various Lightining methods for training and accuracy testing.

First we need to install and import PyTorch Lightining. This can be done with the command

pip install pytorch-lightning

or

conda install -c conda-forge pytorch-lightning

In order for our code to work in Lightning, we need to do the following:

1. Create a subclass of pl.LightningModule and add to it model architecture in __init__ method and forward pass method.
2. Move used optimizer to the configure_optimizers() method.
3. Define the training and validation process in methods training_step and validation_step respectively.
4. (Optional) Implement a testing (test_step method) and prediction process (predict_step method).

It should also be understood that PyTorch Lightning has a built-in translation of models to different devices, depending on where the incoming data from the DataLoaders is located. Therefore, all calls .cuda() or .to(device) should be removed from the code.

Let's also add validation Dataset and DataLoader:

Now our model is ready for training. In Pytorch Lightning, this process is implemented through an object of the Trainer class, which essentially "mixes" the model with any datasets.

Takeaways

• PyTorch allows you to operate on tensors at low level, you have most flexibility.
• There are convenient tools to work with data, such as Datasets and Dataloaders.
• You can define neural network architectures using Sequential syntax, or inheriting a class from torch.nn.Module
• For even simpler approach to defining and training a network - look into PyTorch Lightning