Training and Evaluating a Simple Neural Network in PyTorch

Introduction: The Importance of Evaluation in Neural Networks

Welcome to the first lesson of the "Improving Neural Networks with PyTorch" course. In this course, you will learn practical ways to make your neural networks perform better and avoid common pitfalls. We start with one of the most important steps in any machine learning project: evaluating your model. Evaluation helps you understand how well your model is learning and whether it is likely to perform well on new, unseen data. This is especially important in deep learning, where models can easily become too complex and start to "memorize" the training data — a problem known as overfitting. Overfitting happens when a model learns the training data too well, including its noise and outliers, and as a result, performs poorly on new, unseen data. The model essentially "memorizes" the training set instead of learning general patterns. On the other hand, underfitting occurs when a model is too simple to capture the underlying structure of the data, resulting in poor performance on both the training and validation sets. In this lesson, you will learn how to set up a simple neural network using PyTorch, train it on a dataset, and evaluate its performance using a validation set. This foundation will prepare you for more advanced techniques in later lessons, such as dropout, early stopping, and batch normalization.

Preparing Data for PyTorch

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Building and Training a Simple Neural Network

Now that the data is ready, let's build a simple neural network using PyTorch. We will use a Multi-Layer Perceptron (MLP), which is a basic type of feedforward neural network. In this example, the network has one hidden layer with 64 units and uses the ReLU activation function. The output layer uses a sigmoid activation, which is standard for binary classification tasks.

Let's break down the model architecture and some key design choices:

ReLU Activation Function: The hidden layer uses the ReLU (Rectified Linear Unit) activation function. ReLU introduces non-linearity into the model, which allows it to learn complex patterns in the data. It is also computationally efficient and helps avoid the vanishing gradient problem that can occur with other activation functions like sigmoid or tanh.
Model as a Class Inheriting from nn.Module: The model is defined as a class that inherits from nn.Module. This is the standard way to build models in PyTorch. Inheriting from nn.Module gives you access to useful methods and makes it easy to manage model parameters, move the model between devices (CPU/GPU), and save or load the model.
The forward Method: The forward method defines how the input data flows through the network layers. When you call the model on an input (e.g., model(X_train)), PyTorch automatically calls the forward method to compute the output.

Here is the code for the model:

Python

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(20, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

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

model = MLP()

import torch.nn as nn

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(20, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

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

model = MLP()

To train the model, we need a loss function and an optimizer. For binary classification, we use binary cross-entropy loss (nn.BCELoss). The Adam optimizer is a popular choice for training neural networks, as it adapts the learning rate during training.

Here is how you set up the loss function and optimizer:

Python

import torch.optim as optim

loss_fn = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

import torch.optim as optim

loss_fn = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

The training loop runs for a fixed number of epochs. In each epoch, the model is set to training mode, gradients are cleared, and the model makes predictions on the training data. The loss is computed, gradients are calculated with backward(), and the optimizer updates the model's weights.

Evaluating Model Performance During Training

Evaluating your model as it trains is key to understanding how well it is learning. After each training epoch, we switch the model to evaluation mode and compute the loss on the validation set. This helps you see if the model is improving and whether it might be starting to overfit.

It's important to monitor both the training loss and the validation loss during training. Printing only the validation loss is not sufficient for understanding model performance. By tracking both, you can see if the model is overfitting (training loss decreases while validation loss increases) or underfitting (both losses remain high).

Overfitting: Training loss keeps decreasing, but validation loss starts to increase. The model is learning the training data too well and not generalizing.
Underfitting: Both training and validation loss remain high. The model is too simple or not trained enough to capture the underlying patterns in the data.

In evaluation mode, we use torch.no_grad() to avoid tracking gradients, which saves memory and computation. The training loss and validation loss are printed at each epoch, so you can monitor progress.

Here is the relevant part of the code:

Python

for epoch in range(50):
    model.train()
    optimizer.zero_grad()
    output = model(X_train)
    loss = loss_fn(output, y_train)
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        val_output = model(X_val)
        val_loss = loss_fn(val_output, y_val)
        print(f"Epoch {epoch}: Train Loss = {loss.item():.4f}, Val Loss = {val_loss.item():.4f}")

for epoch in range(50):
    model.train()
    optimizer.zero_grad()
    output = model(X_train)
    loss = loss_fn(output, y_train)
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        val_output = model(X_val)
        val_loss = loss_fn(val_output, y_val)
        print(f"Epoch {epoch}: Train Loss = {loss.item():.4f}, Val Loss = {val_loss.item():.4f}")

The output will look something like this:

Epoch 0: Train Loss = 0.6931, Val Loss = 0.6932
Epoch 1: Train Loss = 0.6900, Val Loss = 0.6901
Epoch 2: Train Loss = 0.6869, Val Loss = 0.6870
...
Epoch 49: Train Loss = 0.5102, Val Loss = 0.5123

Epoch 0: Train Loss = 0.6931, Val Loss = 0.6932
Epoch 1: Train Loss = 0.6900, Val Loss = 0.6901
Epoch 2: Train Loss = 0.6869, Val Loss = 0.6870
...
Epoch 49: Train Loss = 0.5102, Val Loss = 0.5123

A decreasing training and validation loss means the model is learning. If the validation loss starts to increase while the training loss keeps decreasing, it may be a sign of overfitting.

Example: Full Training and Evaluation Script

Let's put everything together. Here is the complete script that prepares the data, builds the model, trains it, and evaluates its performance by printing both the training and validation loss at each epoch:

Python

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

X, y = make_classification(n_samples=1000, n_features=20)
X = StandardScaler().fit_transform(X)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2)

X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).unsqueeze(1)
X_val = torch.tensor(X_val, dtype=torch.float32)
y_val = torch.tensor(y_val, dtype=torch.float32).unsqueeze(1)

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(20, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

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

model = MLP()
loss_fn = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(50):
    model.train()
    optimizer.zero_grad()
    output = model(X_train)
    loss = loss_fn(output, y_train)
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        val_output = model(X_val)
        val_loss = loss_fn(val_output, y_val)
        print(f"Epoch {epoch}: Train Loss = {loss.item():.4f}, Val Loss = {val_loss.item():.4f}")

import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

X, y = make_classification(n_samples=1000, n_features=20)
X = StandardScaler().fit_transform(X)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2)

X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32).unsqueeze(1)
X_val = torch.tensor(X_val, dtype=torch.float32)
y_val = torch.tensor(y_val, dtype=torch.float32).unsqueeze(1)

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(20, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()
        )

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

model = MLP()
loss_fn = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(50):
    model.train()
    optimizer.zero_grad()
    output = model(X_train)
    loss = loss_fn(output, y_train)
    loss.backward()
    optimizer.step()

    model.eval()
    with torch.no_grad():
        val_output = model(X_val)
        val_loss = loss_fn(val_output, y_val)
        print(f"Epoch {epoch}: Train Loss = {loss.item():.4f}, Val Loss = {val_loss.item():.4f}")

When you run this script, you will see both the training loss and validation loss printed for each epoch. This gives you a clear view of how your model is performing as it learns.

Training and Evaluating a Simple Neural Network in PyTorch

Introduction: The Importance of Evaluation in Neural Networks

Preparing Data for PyTorch

Join the 1M+ learners on CodeSignal

Building and Training a Simple Neural Network

Evaluating Model Performance During Training

Example: Full Training and Evaluation Script

Summary and What’s Next