Building Your First Neural Network in PyTorch

# CHAPTER 7

Building Your First Neural Network in PyTorch #

1. Introduction #

2. Learning Objectives #

• Subclass the nn.Module class to create a custom PyTorch model.

• Define network layers (like nn.Linear) inside the _init method.

1. 1. init_(self): This is where you define the layers (the building blocks).

1. 2. forward(self, x): This is where you define *how* the data (x) flows through the layers.

4. Step-by-Step: Building a Model #

import torch
import torch.nn as nn

# 1. Subclass nn.Module
class DigitClassifier(nn.Module):
    def __init__(self):
        super().__init__() # Must call the parent class initializer
        
        # 2. Define the layers
        # Input layer takes 784 pixels. The first hidden layer has 128 neurons.
        self.hidden1 = nn.Linear(in_features=784, out_features=128)
        
        # Second hidden layer takes 128 neurons, outputs 64.
        self.hidden2 = nn.Linear(in_features=128, out_features=64)
        
        # Output layer takes 64 neurons, outputs exactly 10 (one for each digit 0-9)
        self.output = nn.Linear(in_features=64, out_features=10)
        
        # We need activation functions (covered deeply in the next chapter)
        self.relu = nn.ReLU()

    # 3. Define the Forward Pass
    def forward(self, x):
        # Pass the data (x) through hidden layer 1, then apply ReLU
        x = self.hidden1(x)
        x = self.relu(x)
        
        # Pass through hidden layer 2, apply ReLU
        x = self.hidden2(x)
        x = self.relu(x)
        
        # Pass through output layer
        x = self.output(x)
        return x

# Instantiate the brain!
model = DigitClassifier()
print(model)

5. Passing Data Through the Model #

# Create a dummy tensor representing a batch of 1 image, with 784 pixels
dummy_image = torch.randn(1, 784)

# Pass it through the untrained network
predictions = model(dummy_image)

print("Output Shape:", predictions.shape)
# Output Shape: torch.Size([1, 10]) -> 1 image, 10 digit probabilities
print("Raw Predictions:", predictions)

6. Mini Project: The nn.Sequential Shortcut #

import torch.nn as nn

# This creates the exact same architecture as the class above!
sequential_model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 64),
    nn.ReLU(),
    nn.Linear(64, 10)
)

print(sequential_model)

*Why use classes at all then?* Because advanced models (like ResNets or Transformers) don't move in a straight line. Data splits, merges, and skips layers. nn.Sequential cannot handle complex branching logic, but nn.Module classes can do anything.

7. Understanding the Parameters #

# View the random weights of the first layer
print("Weights shape:", model.hidden1.weight.shape)

# Count total parameters in the model
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters to train: {total_params}")
# Output: ~109,386 weights and biases!

8. Common Mistakes #

• Forgetting super()._init(): If you forget this line inside your init method, PyTorch won't know it's supposed to track the weights, and your code will crash immediately.

• Shape mismatch between layers: The outfeatures of Layer 1 MUST match the infeatures of Layer 2. If Layer 1 outputs 128, and Layer 2 expects 100, PyTorch will throw a massive error during the forward pass.

• Define layers, don't define math: Inside init, you define the objects (e.g., self.layer1 = nn.Linear()). Do not put data operations or math here. All math and data routing belongs strictly inside the forward() method.

• Q: Explain the purpose of the init method versus the forward method in a PyTorch nn.Module class.

15. Next Chapter Recommendation #