Building a Neural Network with PyTorch

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:

  1. Setting up the environment
  2. Defining the neural network
  3. Preparing the data
  4. Training the network
  5. 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, 64)
        self.fc2 = nn.Linear(64, 32)
        self.fc3 = nn.Linear(32, 16)
        self.fc4 = nn.Linear(16, 8)
        self.fc5 = nn.Linear(8, 1)
        
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        x = self.fc5(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 pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, r2_score

# Dummy DataFrame for demonstration purposes
data = {
    'feature1': [1, 2, 3, 4, 5],
    'feature2': [5, 4, 3, 2, 1],
    'target': [1, 3, 2, 5, 4]
}
df = pd.DataFrame(data)

# Data preparation
X = df[['feature1', 'feature2']].values
y = df['target'].values

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

# Convert data to PyTorch tensors
X_train_tensor = torch.tensor(X_train, dtype=torch.float32)
X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
y_train_tensor = torch.tensor(y_train, dtype=torch.float32).view(-1, 1)
y_test_tensor = torch.tensor(y_test, dtype=torch.float32).view(-1, 1)

Training the Network

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

def train_nn(model, criterion, optimizer, X_train, y_train, epochs=100, batch_size=2):
    dataset = torch.utils.data.TensorDataset(X_train, y_train)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
    for epoch in range(epochs):
        model.train()
        for batch_X, batch_y in dataloader:
            optimizer.zero_grad()
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()

# Initialize and train the neural network
input_size = X_train.shape[1]
nn_model = NeuralNetwork(input_size)
criterion = nn.MSELoss()
optimizer = optim.Adam(nn_model.parameters(), lr=0.01)

train_nn(nn_model, criterion, optimizer, X_train_tensor, y_train_tensor, epochs=100, batch_size=2)

Evaluating the Network

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

# Evaluate the model
nn_model.eval()
with torch.no_grad():
    y_pred_nn = 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): 3.292425297941236
R2 (Neural Network): nan
/Users/kt536/Library/Python/3.12/lib/python/site-packages/sklearn/metrics/_regression.py:996: UndefinedMetricWarning:

R^2 score is not well-defined with less than two samples.

Visualizing the Network

Finally, let’s visualize the network using torchviz.

from torchviz import make_dot
from IPython.display import Image

dummy_input = torch.randn(1, input_size)
output = nn_model(dummy_input)
vis_graph = make_dot(output, params=dict(nn_model.named_parameters()))
vis_graph.format = 'png'
vis_graph.render('neural_network_visualization')

Image(filename='neural_network_visualization.png')

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.

Happy coding!