Building a Feedforward Neural Network in PyTorch

1. What Are We Building?

A feedforward neural network is organized in layers: an input layer, one or more hidden layers, and an output layer. Information flows in one direction—forward—from the input, through the hidden layers where computation happens, to the output layer that produces the final prediction.

2. The Core Building Block: nn.Module

In PyTorch, every neural network you build is a Python class that inherits from torch.nn.Module. This base class provides all the essential functionality for a neural network, like tracking its layers and parameters.

As the documentation highlights, building a model involves two main steps within your class:

1. __init__() (The Constructor): Here, you define all the layers your network will contain. These are like the Lego bricks of your model. The most common layers for an FNN are:

   • nn.Linear(in_features, out_features): A fully connected layer that applies a linear transformation. in_features is the number of input neurons, and out_features is the number of output neurons.
   • nn.ReLU(): The Rectified Linear Unit activation function. This introduces non-linearity, allowing the network to learn complex patterns.

2. forward(self, x) (The Forward Pass): Here, you define the sequence of operations. You specify the exact path the input data x takes as it flows through the layers you defined in __init__().

3. Building an FNN for Image Classification

Let's make this concrete by building a network to classify handwritten digits from the famous MNIST dataset. This is a classic problem where each input is a 28x28 pixel grayscale image of a digit (0-9).

The process generally follows these steps, which are laid out very clearly in the article "Feedforward Neural Network with PyTorch".

Step 1: Define Hyperparameters and Load Data

Before we build the model, we need to define some key settings. The video below by Patrick Loeber provides a clear, practical walkthrough of setting up for an FNN project.

Step 2: Create the Model Class

Now for the main event. We will create a class named NeuralNet that inherits from nn.Module. It will have one hidden layer.

Here is the code for the model, as shown in the video. Let's break it down:

import torch.nn as nn

class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        # Always call the superclass constructor first
        super(NeuralNet, self).__init__()
        
        # Define the layers
        self.fc1 = nn.Linear(input_size, hidden_size)  # First linear layer (input -> hidden)
        self.relu = nn.ReLU()                          # Activation function
        self.fc2 = nn.Linear(hidden_size, num_classes) # Second linear layer (hidden -> output)
    
    def forward(self, x):
        # Define the data flow
        out = self.fc1(x)       # Pass input through the first linear layer
        out = self.relu(out)    # Apply the ReLU activation
        out = self.fc2(out)     # Pass through the second linear layer
        # No softmax here!
        return out

A Crucial Note: You might wonder why there is no nn.Softmax layer at the end. For classification tasks, the standard loss function is nn.CrossEntropyLoss. This PyTorch function conveniently combines the Softmax activation and the cross-entropy loss calculation. Doing this is more numerically stable than applying them separately. So, the raw output of our final linear layer (called "logits") is exactly what nn.CrossEntropyLoss expects.

Step 3: Instantiate the Model, Loss, and Optimizer

With our class defined, we can now create an instance of our network. We also set up the loss function and the optimizer, connecting all the pieces we've learned.

The code to do this is simple and clean:

# Assuming hyperparameters are defined
# input_size = 784
# hidden_size = 100
# num_classes = 10
# learning_rate = 0.001

model = NeuralNet(input_size, hidden_size, num_classes)

# Loss
criterion = nn.CrossEntropyLoss()

# Optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

At this point, our model is a fully defined neural network, and we have the criterion to measure its error and the optimizer to update its weights. It's ready to be trained!

def __init__(self):
    super(MyNewNet, self).__init__()
    self.fc1 = nn.Linear(50, 128)  # Input is 50, hidden is 128
    self.relu = nn.ReLU()
    self.fc2 = nn.Linear(128, 4)   # Hidden is 128, output is 4

Conclusion

Great work! You've successfully designed the blueprint for a feedforward neural network in PyTorch. You now know how to structure a model as a class, define its layers, and specify the forward pass.

Key Takeaways:

• PyTorch neural networks are built by creating a class that inherits from nn.Module.
• The __init__ method is where you declare the network's layers (e.g., nn.Linear, nn.ReLU).
• The forward method defines the computational graph, specifying how data flows through the layers.
• The dimensions of your nn.Linear layers must match your data's input size, your desired hidden layer sizes, and the number of output classes.
• For multi-class classification, you should output the raw logits from the final linear layer and use nn.CrossEntropyLoss, as it includes the Softmax function internally.

Preview of the Next Lesson:
We have now defined the data loading, the model architecture, the loss function, and the optimizer. All the pieces are on the table. In the next lesson, "Implement a complete training loop in PyTorch," we will assemble them all into a single script to train our network, watch it learn, and evaluate its performance on the test data.