Learning PyTorch: Neural Networks

Python

Author

Kai Tan

Published

May 1, 2024

Introduction

In this blog post, we will explore how to build a neural network using PyTorch. PyTorch is a powerful and flexible deep learning framework that makes it easy to define, train, and evaluate neural networks. We will cover the following steps:

Setting up the environment
Defining the neural network
Preparing the data
Training the network
Evaluating the network

Setting Up the Environment

First, let’s ensure that we have PyTorch installed. You can install PyTorch using pip:

pip install torch torchvision

Defining the Neural Network

Next, we will define a simple neural network with five layers. Each hidden layer will use the ReLU activation function.

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

class NeuralNetwork(nn.Module):
    def __init__(self, input_size):
        super(NeuralNetwork, self).__init__()
        self.fc1 = nn.Linear(input_size, 16)
        self.fc2 = nn.Linear(16, 8)
        self.fc3 = nn.Linear(8, 1)
        
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

Preparing the Data

For this example, we will create a simple dataset. In a real-world scenario, you would use a dataset from a file or an online source.

import torch
import numpy as np
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score
from sklearn.preprocessing import StandardScaler

# Set random seed for reproducibility
seed = 42
torch.manual_seed(seed)
np.random.seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load the California housing dataset (first 100 samples for simplicity)
california = fetch_california_housing()
X, y = california.data[:100], california.target[:100]

# Split into train/test sets
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=seed
)

# Standardize features
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# Convert data to PyTorch tensors and move to device
X_train_tensor = torch.from_numpy(X_train).float().to(device)
X_test_tensor = torch.from_numpy(X_test).float().to(device)
y_train_tensor = torch.from_numpy(y_train).float().unsqueeze(1).to(device)
y_test_tensor = torch.from_numpy(y_test).float().unsqueeze(1).to(device)

# Create a DataLoader for batching
dataset = torch.utils.data.TensorDataset(X_train_tensor, y_train_tensor)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=20, shuffle=True)

Using device: cpu

Training the Network

Now, let’s define the training function and train our neural network.

# Train the neural network model
input_size = X_train.shape[1]
model = NeuralNetwork(input_size)
loss_fn = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)
n_epochs = 100

for epoch in range(n_epochs):
    model.train()
    for batch_X, batch_y in dataloader:
        # step 1: make prediction (forward pass)
        outputs = model(batch_X)
        # step 2: compute loss
        loss = loss_fn(outputs, batch_y)
        # step 3: clear old and compute current gradients
        optimizer.zero_grad() # clear old gradients
        loss.backward()       # compute current gradients
        # step 4: update parameters
        optimizer.step()

Evaluating the trained Network

After training, we need to evaluate our network’s performance on the test set.

# Evaluate the model
model.eval()
with torch.no_grad():
    y_pred_nn = model(X_test_tensor).numpy()

mse_nn = mean_squared_error(y_test, y_pred_nn)
r2_nn = r2_score(y_test, y_pred_nn)

print(f'MSE (Neural Network): {mse_nn}')
print(f'R2 (Neural Network): {r2_nn}')

MSE (Neural Network): 0.143800164172336
R2 (Neural Network): 0.8298946833143164

Conclusion

In this blog post, we covered the basics of building a neural network using PyTorch. We went through setting up the environment, defining the neural network, preparing the data, training the network, and evaluating its performance. PyTorch provides a flexible and intuitive framework for developing deep learning models, making it a popular choice for researchers and practitioners alike.

Feel free to experiment with different architectures, hyperparameters, and datasets to further enhance your understanding and build more complex models.

Online Resources on PyTorch and Deep Learning