How to train you first model using pytorch?

Hello Guys ,

Welcome to new post , today we are going to learn about hands on example of building simple neural network using most popular deep learning framework pytorch.

Photo by Content Pixie on Unsplash

Define Dataset

Typically in any machine learning model training process we pass some X input and expect model to predict Y as output. Lets follow similar pattern using below sample data. We have series of input [1,2,3] and trying to predict target [4] as Y.

Lets create a simple dataset and cast them as PyTorch tensor for model training. Below is sample code where dimension of x is 3 * 3 as input and expecting corresponding Y label.

x_raw = [[1,2,3],
         [2,3,4],
         [3,4,5]]
y_raw = [4,5,6]
X = torch.tensor(x_raw, dtype=torch.float32)
y = torch.tensor(y_raw, dtype=torch.float32)

Great now are input data is ready now next step is do design neural network to process input data to make final prediction. These types of problems are called regression problem where you are not predicting pre-defined classes instead trying to predict number which will vary.

But we know that we are passing 3 input information and expecting 1 target output. In this case my simplest network would look like below.

Now since our data is ready , let start implementing SimpleNN class by inheriting nn.Module. nn.Module provides useful methods and attributes which can be use for model training.

import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.linear = nn.Linear(3, 1)  # 3 input features, 1 output feature

    def forward(self, x):
        return self.linear(x)

# Create an instance of our network
model = SimpleNet()

In above code , I am using nn.Linear which is linear transformation layer . It applies below linear formula to calculate target Y. https://pytorch.org/docs/stable/generated/torch.nn.Linear.html

Linear Transformation

Also not to forget about forward function , it is automatically called during model training process and its useful for froward pass during model training.

You can print your model to verify your model architecture.

print(model)
#SimpleNet(
#  (linear): Linear(in_features=3, out_features=1, bias=True)
#)

Training Model

Alright now lets define other parameters which are essential for model training. We need two critical element for model training i.e. Loss function and Optimizer.

We regression problem we can use Mean Squared Error and for optimizer we can use Stochastic Gradient Descent (SGD). If you need more details on available loss function you can refer pytorch documentations.

# Loss and Optimizer
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

We are almost set to start training model now. In below code we have initialized number of epochs to 100 but it could vary based on dataset size.

number_of_epochs = 100
for epoch in range(number_of_epochs):  # number of epochs

    pred_y = model(X) # this is where forward pass happens
    loss = criterion(pred_y, y) # calculate loss 
    
    # Zero gradients, perform a backward pass, and update the weights.
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 10 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item()}')

After 10 epochs , it will print current loss and we can observe that after each pass our loss is decreasing which is good sing that our model is learning. If your loss is not reducing that it becomes good point to stop model training process.

Making Predictions

Now once all epochs are finished , you can now use model to make predictions. From below code we can see that we I pass unseen data to model it predicted number close to 7 which is what I expect model to do.

test_x = torch.tensor([4,5,6], dtype=torch.float32)
y = model(test_x)
print(y)

# tensor([7.1486], grad_fn=<AddBackward0>)

Saving Your Pre-Trained Model

Once you finish your model training process , its good idea to save your model otherwise you might loose your model object and you will to re-train mode.

torch.save(model, 'model.pth')

And there you have it! A simple example of building and training a neural network using PyTorch. This example demonstrates the fundamental steps involved in setting up a neural network, defining a loss function, and training the model using an optimizer. While this network is quite basic, the principles apply to more complex models and datasets.

Complete Code

import torch
import torch.nn as nn
import torch.optim as optim
import matplotlib.pyplot as plt

x_raw = [[1,2,3],
         [2,3,4],
         [3,4,5]]
y_raw = [4,5,6]

X = torch.tensor(x_raw, dtype=torch.float32)
y = torch.tensor(y_raw, dtype=torch.float32)


class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.linear = nn.Linear(3, 1)  # 3 input features, 1 output feature

    def forward(self, x):
        return self.linear(x)

# Create an instance of our network
model = SimpleNet()
print(model)

# Loss and Optimizer
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

number_of_epochs = 100
for epoch in range(number_of_epochs):  # number of epochs

    pred_y = model(X) # this is where forward pass happens
    loss = criterion(pred_y, y) # calculate loss 
    
    # Zero gradients, perform a backward pass, and update the weights.
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch % 10 == 0:
        print(f'Epoch {epoch}, Loss: {loss.item()}')

# make prediction
print("====Make prediction ====")
test_x = torch.tensor([4,5,6], dtype=torch.float32)
y = model(test_x)
print(y)

# save model
torch.save(model, 'model.pth')

Keep Learning !! Keep Sharing !!

Loss Functions — https://pytorch.org/docs/stable/nn.html#loss-functions