PyTorch Tutorial – Build the Model

In this tutorial, we would now have to create a neural network model. To do this in Pytorch we do the following:

• create a class that extends the nn.Module class
•  define the layers of the network in the init method
• specify how data would pass through the network in the forward function
• define a device which could be a gpu, cpu or mps

1. Declare the Device
2. Create the Neural Network Model
3. Understanding the Model

Declare the Device

Here we use GPU if it exists, else we use the CPU.

# Declare the device: cuda or mps or cpu
device = (
    "cuda"
    if torch.cuda.is_available()
    else "mps"
    if torch.backends.mps.is_available()
    else "cpu"
)
print(f"Using device {device}")

2. Create the Neural Network Model

• create a class that extends the nn.Module class
• define the layers of the network in the init method
• specify how data would pass through the network in the forward function

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        
        # create the input layer
        self.flatten = nn.Flatten()
        
        # create the 2 hidden layers
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU(),
            nn.Linear(512, 512),
            nn.ReLU(),
            nn.Linear(512, 10)
        )
        
    def forward(self, x):
        x = self.flatten(x)
        logits = self.linear_relu_stack(x)
        return logits
    
model = NeuralNetwork().to(device)
print(model)

3. Understanding the Model

Now I will take some time to explain the model we created.

nn.Flatten() – this layer is used to flatten the input tensor from a 2D image (28 by 28 pixels) to a 1-d tensor (size 784) before feeding it into the fully connected layers.

nn.Sequential() – this is used to create a sequence of layers.

nn.ReLU() – this applies the activation function (Rectified Linear Unit) which introduces non-linearity to the model by zeroing out negative values

In the forward() function, we do the following:

• flattening – the input x, which is a batch of 2d images is flattened to a 1d tensor
• the flattened input layer is passed through the sequence of layers to produce the output logits (