Training a Convent from Scratch

PyTorch - Training a Neural Network from Scratch

In this chapter, we will focus on creating and training a simple feed-forward neural network from scratch using PyTorch’s nn.Module. This involves defining the network architecture, the forward pass, and then implementing a training loop.

Step 1: Define the Neural Network

We’ll create a class that inherits from torch.nn.Module. Inside, we define the layers of our network. For a simple network, this might include linear layers and activation functions.

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleNeuralNetwork(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNeuralNetwork, self).__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.output_size = output_size

        # Define layers
        self.fc1 = nn.Linear(self.input_size, self.hidden_size) # Input layer to hidden layer
        self.relu = nn.ReLU() # Activation function for hidden layer
        self.fc2 = nn.Linear(self.hidden_size, self.output_size) # Hidden layer to output layer
        self.sigmoid = nn.Sigmoid() # Activation function for output layer (e.g., for binary classification)

    def forward(self, x):
        # Define the forward pass
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.sigmoid(out) # Use sigmoid if output is a probability (0-1)
        return out

# Example instantiation:
# input_features = 10
# hidden_units = 20
# output_classes = 1 (for binary classification)
# model = SimpleNeuralNetwork(input_features, hidden_units, output_classes)

Key changes from older approaches:

We use nn.Linear for fully connected layers, which handles weight and bias initialization internally.
We use standard activation functions like nn.ReLU and nn.Sigmoid from torch.nn.
The forward method defines how data flows through the network.

Step 2: Prepare Data, Loss Function, and Optimizer

Before training, you need data, a loss function to measure error, and an optimizer to update model weights.

# Example: Dummy data for illustration
# Assume X_train is your input data and y_train are your labels
# For a network with input_size=2, output_size=1:
X_train = torch.randn(100, 2)  # 100 samples, 2 features each
y_train = torch.randint(0, 2, (100, 1)).float() # 100 binary labels

# Instantiate the model
input_dim = 2
hidden_dim = 3
output_dim = 1
model = SimpleNeuralNetwork(input_dim, hidden_dim, output_dim)

# Loss Function
# For binary classification with sigmoid output, BCELoss is common.
# If it were a regression task, MSELoss might be used.
criterion = nn.BCELoss()

# Optimizer
# Adam and SGD are popular choices.
learning_rate = 0.01
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

print("Model Architecture:")
print(model)

PyTorch’s autograd system automatically computes gradients. We no longer need to manually implement backward pass logic or derivative functions like sigmoidPrime.

Step 3: Training Loop

The training loop iterates over the data multiple times (epochs), performing forward and backward passes.

num_epochs = 100

for epoch in range(num_epochs):
    # Forward pass: Compute predicted y by passing X to the model
    outputs = model(X_train)
    loss = criterion(outputs, y_train)

    # Backward pass and optimize
    optimizer.zero_grad()  # Zero the gradients before running the backward pass.
    loss.backward()        # Perform backpropagation: compute gradients of the loss w.r.t. model parameters.
    optimizer.step()       # Calling step() causes the optimizer to update the parameters.

    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

print("\nTraining finished.")

Step 4: Prediction and Saving/Loading Weights

After training, you can use the model for predictions and save its state.

# Prediction
def predict(model, X_new):
    model.eval()  # Set the model to evaluation mode (important for layers like Dropout, BatchNorm)
    with torch.no_grad(): # Disable gradient calculations for inference
        predictions = model(X_new)
        predicted_classes = (predictions > 0.5).float() # Assuming binary classification with 0.5 threshold
    return predicted_classes

# Example prediction with new data
X_test = torch.randn(5, 2) # 5 new samples
predictions = predict(model, X_test)
print("\nPredicted data based on trained weights: ")
print("Input (scaled): \n" + str(X_test))
print("Output (predicted classes): \n" + str(predictions))

# Saving and Loading Model Weights
MODEL_PATH = "simple_nn.pth"

# Save model state_dict (recommended)
torch.save(model.state_dict(), MODEL_PATH)
print(f'\nModel state_dict saved to {MODEL_PATH}')

# To load the model weights:
# First, instantiate the model again with the same architecture
loaded_model = SimpleNeuralNetwork(input_dim, hidden_dim, output_dim)
loaded_model.load_state_dict(torch.load(MODEL_PATH))
loaded_model.eval() # Set to evaluation mode after loading
print('Model loaded successfully.')

# You can also save the entire model (less flexible):
# torch.save(model, "entire_model.pth")
# loaded_model = torch.load("entire_model.pth")

Using model.eval() is crucial during inference or validation to ensure layers like dropout or batch normalization behave correctly. torch.no_grad() disables gradient computation, saving memory and speeding up inference.