How to Cook Neural Nets with PyTorch
A recipe for deep learning projects using PyTorch
The PyTorch Training Recipe
Ok, so you’ve decided on the dish (your neural network) and now you need to cook (train) it using PyTorch. But wait, this is not a simple ‘fast food’ task like running some variation of fit() and then eval() method unlike other popular python machine learning libraries (e.g. Scikit Learn, Keras). You’ll need to follow a recipe (process) and define these methods yourself. Why, you may ask? Well, there is certainly a tradeoff here. While PyTorch is known for its interpretability, it’s pythonic and debuggable code and the ability to create dynamic computation graphs, it is not as practical as Keras because it lacks a higher level API, especially for training and evaluating models. Again, it’s a tradeoff: flexibility vs practicality. PyTorch focuses on the former, but there are libraries that make training and evaluating easier, like ignite, fastai and lightning (to name a few).
So back to our recipe… We’ll take a look at the big picture first and then zoom in into each step. Let’s take a step back and look at what a typical ML process may look like. It turns out there’s a simple recipe you can follow. It goes something like this:
- Load and Preprocess the Data
- Define Model, Optimizer and Loss
- Train and Validate the Model
- Evaluate the Model
Now let’s take a look at how one might do this in PyTorch. I’ll show you a recipe with some boilerplate code (no pun intended =) that you can follow to develop and train your own model. Then, it’ll be up to you to change the ingredients and spice up the recipe here and there!
1. Loading and Preprocessing the Data
Define Transforms (for Image Classification only)
If you’re doing image classification, you’ll likely need to transform your data. PyTorch’s torchvision module has a nice submodule called transforms that lets you compose a pipeline of transformations. Unlike in Keras, here you get to define the order of the transformations. These transformations can include image augmentation techniques, like resize, crop, color jitter, image flip etc. A detailed list of the transformations can be found here on the PyTorch website.
If your dataset is nice and standardized (all images have the same size, are pretty much centered, no color variations, images are not rotated or flipped) like MNIST, then you probably just need to transform your images to tensors and normalize them. Below is an example of a simple transformation pipeline you can use for standardized datasets using the transforms module.
In the example above, the ToTensor() transform takes as input an image and then scales each pixel value between 0 and 1 by dividing the value by 255 for each channel (we are working with grayscale images above, which have only one channel). Next, the tensor is normalized so that the values are between -1 and 1 (previously scaled pixels are subtracted from the mean of 0.5 then divided by the standard deviation of 0.5). This works like a pipeline and is super convenient for preprocessing images.
If however, you are trying to classify images of cats and dogs or any real world dataset, you will likely benefit from data augmentation. The main idea is that your model will generalize better when it is trained on a greater variety of data. Data augmentation is a technique to leverage your existing data by making variations (or transformations) to the original content. Here’s an example of how one might approach this task:
The names of the transformers are pretty intuitive and you can check out the definition of these and others by browsing the PyTorch documentation. Notice how you can define different transforms for your train and test sets. Also, notice how we now have three channels instead of one in the normalization step.
Download the datasets
PyTorch provides some options for you to include or download a dataset. The torchvision package contains preloaded datasets, like CIFAR10 and MNIST, which can be very handy if you want to benchmark your model or just want to learn image classification. If you want to classify your own images or download them from somewhere on the web, you’ll need another way to load the dataset. For that, PyTorch provides the ImageFolder class, which just assumes that you have all your images stored into directories that are named after each of the class labels (i.e. the folder ‘dog’ will contain images of dogs, the folder ‘cat’ will contain images of cats). You should have a train and a test folder with the same structure. That way, you can load your data like this:
Now let’s say that you need to do some cleaning and preprocessing on your images or if you’re not doing image classification and rather need to write a text based classifier that requires the text to be preprocessed or you need to predict sales on a tabular dataset that requires cleaning or some sort of data wrangling before it’s fed into the model. For these cases, you can create your own dataset by extending the Dataset class from torch.utils.data. You’ll need to define three methods, following a template like this one:
This is a more elegant and flexible approach to define a Dataset, and it’s also generic, it applies to any type of dataset. Once you define your dataset, you can then instantiate the class and pass in the parameters for multiple sets, like train and test. So you’d have something in the lines of…
You can get more information about creating custom classes as well as other customized data related classes here.
Create a Validation Set
Ok, it’s time to split and create a validation set to help you achieve your goal. Note this is an optional step, but definitely a best practice and highly recommended for evaluating and optimizing your model. One way to create a validation set is to use the random_split function from torch.utils.data. It’s pretty straightforward and only requires the dataset you want to split and the size of each split.
Define the DataLoaders
Once you’ve defined your datasets, you’ll need to wrap them in a class called DataLoader. This class basically is an easy way to let you load batches of data and includes parallel processing for free without needing to do any extra coding. Pretty cool stuff. Let’s check it out:
2. Define Model, Optimizer and Loss
Great, you’re now ready to start defining your model. There are a few different ways to define a model and I cover that in another article. But for now, we’ll just pick the class approach as it’s one of the most popular approaches used in PyTorch and it’s pretty intuitive. This is where you can get creative. Depending on your problem, you can try different types of neural networks and architectures. For illustration purposes, I’ve chosen to implement a simple CNN.
You can find more info about creating your custom models here. Alright, now that you’ve defined your model you can instantiate it and define the optimizer and loss function. These are all one-liners and pretty easy to implement. Let’s instantiate the model first.
Output:
MyCNN(
(fc1): Linear(in_features=32, out_features=16, bias=True)
(fc2): Linear(in_features=16, out_features=8, bias=True)
(fc3): Linear(in_features=8, out_features=2, bias=True)
)Next, to define the optimizer you’ll need to import the optim module from PyTorch. PyTorch has many optimizers, and you can lookup the full list here to see which one to use. Each optimizer has different parameters, but most require at least a learning rate, referred to as lr. In order for PyTorch to optimize your network, you’ll first need to pass in the model parameters. It should look something like this:
You’ll then need to define the loss function that PyTorch will use to know how far your model predictions are from the actual targets. The choice of the loss function is dependent on the type of problem you’re trying to solve, and you can find a more detailed list of losses here. For now, we’ll use the Binary Cross Entropy loss, suited for binary classification problems. We’ll call that loss criterion, like below:
3. Train and Validate the Model
Ok, now you’re getting into the real cooking. The train loop is not as straight forward if compared to sklearn or keras as I mentioned previously, but it is very intuitive and can help you understand what is going on under the hood, plus you can debug it very easily by inserting a trace. You can basically perform a train loop just a few steps…
What you saw above is a bare bone PyTorch training loop with validation. In practice, however, you’ll probably want to accumulate the accuracy or your metric of choice and/or apply early stopping at the end of each epoch. The good thing is that this is fully customizable. Really, the sky is the limit. Though for the purpose of this article, we’ll keep it simple.
4. Evaluate the Model
After you’ve trained your model, the final step is to predict and evaluate the results. The test loop is very similar to the validation loop, but instead, you’ll query the test_loader to read in the test data and only run through the test data once to evaluate the model. In other words, you won’t need to worry about epochs, since you just want to see how the model performs once on each instance.
Closing Remarks
Now that you’ve seen how simple it is, try cooking up a Neural Network with PyTorch yourself! The recipe is basically the same, but you can change a few ingredients here and there to customize it to your taste buds. There are other ways to approach these steps, but I decided to show you the one I consider the most straightforward. Let me know what you think, I’d love to know how you approach the recipe!
