Deep Learning in PyTorch with CIFAR-10 dataset

Sergio Alves
2 min readJun 12, 2020

--

In this post, we will learn how to build a deep learning model in PyTorch by using the CIFAR-10 dataset.

PyTorch

PyTorch is a Machine Learning Library created by Facebook. It works with tensors, which can be defined as a n-dimension matrix from which you can perform mathematical operations and build Deep Learning Models.

Deep Learning

This subfield of AI seeks to emulate the learning approach that humans use to obtain certain types of knowledge. In its simplest form, deep learning can be seen as a way to automate predictive analytics.

CIFAR-10 Dataset

The CIFAR-10 dataset consists of 60000 32x32 colour images in 10 classes, with 6000 images per class. There are 50000 training images and 10000 test images.

You can find more information about CIFAR-10 dataset from here.

Deep Learning Model Implementation

For the implementation of this deep learning model, we will go through the following steps:

  1. Import libraries
  2. Preparing the data
  3. Model
  4. Using a GPU
  5. Training the model

Import libraries

Preparing the Data

Here, we imported the datasets and converted the images into PyTorch tensors.

By using the classes method, we can get the image classes from the dataset.

With this for loop, we can get the number of images per class. It goes through all the dataset, add the class name to a dictionary if it doesn’t exist there yet and counts each image per class.

Now, we’ll split the dataset into two groups: training and validation datasets.

We used a validation set with 5000 images (10% of the dataset). To ensure we get the same validation set each time, we set PyTorch’s random number generator to a seed value of 43.

Here, we used the random_split method to create the training and validations sets.

We created dataloaders for training, validation and test sets. We set shuffle=True for the training dataloader, so that the batches generated in each epoch are different, and this randomization helps generalize & speed up the training process. On the other hand, since the validation dataloader is used only for evaluating the model, there is no need to shuffle the images.

Also, we set pin_memory=True because we will push the data from the CPU into the GPU and this parameter lets theDataLoader allocate the samples in page-locked memory, which speeds-up the transfer.

Here, we can visualize a batch of data using the make_grid helper function from Torchvision.

Model

Using a GPU

Training the Model

Plot the losses and the accuracies to check if you’re starting to hit the limits of how well your model can perform on this dataset.

Finally, evaluate the model on the test dataset report its final performance.

--

--