Create a Neural Network in Python with PyTorch

Neural networks are at the heart of deep learning and have been widely used in fields like image recognition, natural language processing, and more. This guide walks you through creating a basic neural network in Python using the powerful PyTorch library.

Step 1: Install Required Libraries

Before starting, ensure you have Python installed on your machine. Then, install PyTorch by following the instructions on the PyTorch official website. Additionally, install NumPy for numerical computations.

pip install torch torchvision numpy

Step 2: Import Libraries

Start by importing the necessary libraries.

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

Step 3: Define the Dataset

For this tutorial, we'll use a simple dataset to perform regression (predicting outputs based on inputs).

# Generate sample data
np.random.seed(42)
x = np.random.rand(100, 1).astype(np.float32)  # Input features
y = 3 * x + 2 + 0.1 * np.random.randn(100, 1).astype(np.float32)  # Target

# Convert to PyTorch tensors
x_tensor = torch.from_numpy(x)
y_tensor = torch.from_numpy(y)

Step 4: Define the Neural Network

Create a simple feedforward neural network with one hidden layer.

# Define the model
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.hidden = nn.Linear(1, 10)  # Hidden layer with 10 units
        self.output = nn.Linear(10, 1)  # Output layer

    def forward(self, x):
        x = torch.relu(self.hidden(x))  # Apply ReLU activation
        x = self.output(x)              # Final output
        return x

# Initialize the model
model = SimpleNN()

Step 5: Define the Loss Function and Optimizer

The loss function quantifies the error in predictions, and the optimizer updates the model's weights based on the error.

# Mean Squared Error Loss
loss_function = nn.MSELoss()

# Stochastic Gradient Descent Optimizer
optimizer = optim.SGD(model.parameters(), lr=0.01)

Step 6: Train the Neural Network

Training involves feeding data into the model, calculating the loss, and updating the weights to minimize the loss.

# Number of epochs (iterations)
epochs = 1000

for epoch in range(epochs):
    # Forward pass
    predictions = model(x_tensor)
    loss = loss_function(predictions, y_tensor)

    # Backward pass
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # Print loss every 100 epochs
    if (epoch + 1) % 100 == 0:
        print(f'Epoch {epoch+1}/{epochs}, Loss: {loss.item():.4f}')

Step 7: Evaluate the Model

After training, you can evaluate how well the model predicts on the input data.

# Switch model to evaluation mode
model.eval()

# Make predictions
with torch.no_grad():
    predictions = model(x_tensor)

# Convert predictions to NumPy for visualization
predictions = predictions.numpy()

# Visualize the results
import matplotlib.pyplot as plt
plt.scatter(x, y, label='True Data')
plt.plot(x, predictions, color='red', label='Model Predictions')
plt.legend()
plt.show()

Step 8: Save and Load the Model

Save your trained model and load it later for inference.

# Save the model
torch.save(model.state_dict(), 'simple_nn.pth')

# Load the model
model.load_state_dict(torch.load('simple_nn.pth'))
model.eval()