PyTorch - Linear Regression

Linear regression is a fundamental supervised machine learning algorithm used to model the linear relationship between a dependent variable (target) and one or more independent variables (predictors or features). In its simplest form (simple linear regression), we try to find the best-fitting straight line through the data points.

The relationship is modeled by the equation:

Y = wX + b

Where Y is the predicted target, X is the input feature, w is the weight (slope of the line), and b is the bias (y-intercept). Our goal is to find the optimal values for w and b that minimize the difference between the predicted Y and the actual target values.

This involves two key concepts:

Cost Function: A function that measures how well our model is performing. For linear regression, Mean Squared Error (MSE) is commonly used.
Gradient Descent: An optimization algorithm used to minimize the cost function by iteratively adjusting the model parameters (w and b) in the direction opposite to the gradient of the cost function.

Interpreting the Result

For the equation Y = wX + b:

The value of ‘w’ (weight or slope) indicates how much Y changes for a one-unit change in X.
The value of ‘b’ (bias or y-intercept) is the value of Y when X is 0.
The goal of training is to find ‘w’ and ‘b’ such that the line Y = wX + b best fits the training data.

Graphically, linear regression tries to find a straight line that passes as closely as possible to all the data points plotted on a 2D plane (for simple linear regression).

Let’s implement linear regression using PyTorch:

Step 1: Import Necessary Libraries

We’ll need torch for PyTorch functionalities, numpy for creating sample data, and matplotlib for plotting.

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
# For a slightly nicer plot style (optional)
# import seaborn as sns
# sns.set_style(style='whitegrid')
# plt.rcParams["patch.force_edgecolor"] = True

Step 2: Prepare Training Data

Let’s create some synthetic data with a known linear relationship and add some noise.

# Define true parameters
true_w = 2.0
true_b = 3.0
num_samples = 256

# Generate random x values
X_numpy = np.random.rand(num_samples, 1).astype(np.float32)

# Generate y values based on y = w*x + b + noise
noise = np.random.randn(num_samples, 1).astype(np.float32) * 0.2 # Smaller noise
y_numpy = X_numpy * true_w + true_b + noise

# Convert NumPy arrays to PyTorch tensors
X_tensor = torch.from_numpy(X_numpy)
y_tensor = torch.from_numpy(y_numpy)

# Optional: Plot the raw data
# plt.figure(figsize=(8, 5))
# plt.scatter(X_numpy, y_numpy, label='Original data', s=10)
# plt.xlabel('x')
# plt.ylabel('y')
# plt.legend()
# plt.show()

Step 3: Implement Linear Regression with PyTorch

We’ll define a simple linear regression model using torch.nn.Module.

# Define the model
class LinearRegressionModel(nn.Module):
    def __init__(self, input_dim, output_dim):
        super().__init__() # Modern Python 3 super() call
        self.linear = nn.Linear(input_dim, output_dim)

    def forward(self, x):
        out = self.linear(x)
        return out

# Model parameters
input_dim = 1 # Single feature
output_dim = 1 # Single output value

# Instantiate the model
model = LinearRegressionModel(input_dim, output_dim)

# Define loss function (Mean Squared Error)
criterion = nn.MSELoss()

# Define optimizer (Stochastic Gradient Descent)
learning_rate = 0.1
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

# Training loop
epochs = 100
for epoch in range(epochs):
    # Forward pass: Compute predicted y by passing x to the model
    predictions = model(X_tensor)

    # Compute loss
    loss = criterion(predictions, y_tensor)

    # Backward pass: Zero gradients, perform backpropagation, and update weights
    optimizer.zero_grad()  # Clear previous gradients
    loss.backward()        # Compute gradients of loss w.r.t. model parameters
    optimizer.step()       # Update model parameters

    if (epoch + 1) % 10 == 0:
        # Get current learned parameters (w, b)
        # .item() gets the value as a Python number from a single-element tensor
        # .data gives access to the tensor without gradient tracking
        learned_w = model.linear.weight.data.item()
        learned_b = model.linear.bias.data.item()
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}, Learned w: {learned_w:.3f}, Learned b: {learned_b:.3f}')

# Plot the results after training
def plot_fit(title="Fitted Line"):
    plt.figure(figsize=(10, 6))
    plt.title(title)
    plt.scatter(X_numpy, y_numpy, s=15, label='Original data')

    # Get predictions from the trained model
    with torch.no_grad(): # Disable gradient calculation for inference
        predicted_values = model(X_tensor).numpy()

    plt.plot(X_numpy, predicted_values, color='red', label='Fitted line')
    plt.xlabel('x (input)')
    plt.ylabel('y (target/predicted)')
    plt.legend()
    plt.grid(True)
    plt.show()

plot_fit('Linear Regression Fit After Training')

# You can also inspect the final learned parameters:
final_w = model.linear.weight.item()
final_b = model.linear.bias.item()
print(f'\nTrue w: {true_w}, True b: {true_b}')
print(f'Learned w: {final_w:.4f}, Learned b: {final_b:.4f}')

The output will first print the loss and learned parameters at intervals during training. Finally, it will show a plot with the original data points and the regression line fitted by the PyTorch model. The learned ‘w’ and ‘b’ should be close to the true values true_w and true_b used to generate the data.