Training in Pytorch

Build and train a neural network with Pytorch to classify real-life objects.

Now that we have mastered optimization and training, let’s get our hands dirty and try to train an NN in Pytorch.

We will build a network with an input of size 3072, 3 linear layers with dimensions 128, 64, 10, and 2 Relu layers in between.

For our training data, we will use a well-known image dataset called CIFAR. CIFAR includes 32*32 RGB (3 channels) images with real-world objects. Our goal is to classify the image into 10 different classes. This is why the input features will be 32∗32∗3, and the output will be the number of classes, which is 1010.

Here is an illustration of the dataset:

CIFAR10 dataset
CIFAR10 dataset


We will also use the Vanilla SGD algorithm torch.optim.SGD and the Cross-Entropy Loss torch.nn.CrossEntropyLoss.

You will be provided with the data loading and transformation functions.

Here is what you need to do:

  • Use the provided functions to load the data in Pytorch.
  • Build the model.
  • Train it for 1 epoch in the entire dataset.
  • Return the final loss.

How to use the optimizer and gradients

Given that you have a neural network called net and a data loader that gives you a bunch of data, the training functionality will look something like this:

inputs, labels = data
#zero the parameter gradients

#forward pass
outputs = net(inputs)
#loss calculation
loss = criterion(outputs, labels)
#backward pass - calculate the gradients 
#apply the gradients using the update rule

Of course, you have to loop over the training data of your data loader.


If you have any questions or you are not sure how to proceed, feel free to check the original Pytorch documentation and look at some examples. You can find the links in the Appendix.

Get hands-on with 1200+ tech skills courses.