Linear Regression with PyTorch

What is Linear Regression ?

Linear Regression is an approach that tries to find a linear relationship between a dependent variable and an independent variable by minimizing the distance as shown below.

Import necessary libraries

# Import PyTorch and Matplotlib
import torch
from torch import nn
import matplotlib.pyplot as plt

After that let's create device-agnostic code

device = "cuda"if torch.cuda.is_available() else"cpu"

Getting Data

Create some data using the linear regression formula of y = weight * X + bias

weight = 0.7

bias = 0.3

Create range values

start = 0

end = 1

step = 0.02

Create X and y (features and labels)

X = torch.arange(start, end, step).unsqueeze(dim=1)

y = weight * X + bias

Split data

train_split = int(0.8 * len(X))

X_train, y_train = X[:train_split], y[:train_split]

X_test, y_test = X[train_split:], y[train_split:]

Plot Data

def plot_predictions(train_data, train_labels, test_data, test_labels, predictions=None):
# Plot training data in blue
plt.scatter(train_data, train_labels, c="b", label="Training Data")

# Plot test data in green
plt.scatter(test_data, test_labels, c="g", label="Test Data")

# Are there any predictions
if predictions is not None:
# Plot the predictions
plt.scatter(test_data, predictions, c="r", label="Predictions")

plt.legend()


# call function to plot data
plot_predictions(X_train, y_train, X_test, y_test)

Building a PyTorch Linear Model

Create our linear model
class LinearRegressionModel(nn.Module):
def __init__(self):
super().__init__()
# Use nn.Linear() for creating the model parameters
self.linear_layer = nn.Linear(in_features=1, out_features=1)

def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.linear_layer(x)

# Set manual seed
torch.manual_seed(42)
model = LinearRegressionModel()
model, model.state_dict()

Training

For training we need:

• Loss Function
• Optimizer
• Training Loop
• Testing Loop

# Setup Loss Function
loss_fn = nn.L1Loss() # similar to MAE

# Setup our Optimizer
optimizer = torch.optim.SGD(params=model.parameters(), lr=0.01)

# Let's write a training loop
torch.manual_seed(42)

epochs = 200
for epoch in range(epochs):
model.train()

# 1. Forward pass
y_pred = model(X_train)

# 2. Calculate the loss
loss = loss_fn(y_pred, y_train)

# 3. Optimizer zero grad
optimizer.zero_grad()

# 4. Perform backpropagation
loss.backward()

# 5. Optimizer step
optimizer.step()

### Testing
model.eval()
with torch.inference_mode():
test_pred = model(X_test)
test_loss = loss_fn(test_pred, y_test)

# Print out what's happening
if epoch % 10 == 0:
print(f"Epoch {epoch} \t Loss: {loss} \t Test Loss {test_loss}")

Making evaluation and prediction

# Turn model into evaluatin mode
model.eval()

# Make predictions on the test data
with torch.inference_mode():
y_preds = model(X_test)

y_preds

# Visualize predictions
plot_predictions(X_train, y_train, X_test, y_test, predictions=y_preds)

Saving & Loading a trained model

from pathlib import Path

# 1. Create a model directory
MODEL_PATH = Path("models")
MODEL_PATH.mkdir(parents=True, exist_ok=True)

# 2. Create model save path
MODEL_NAME = "pytorch_model0.pth"
MODEL_SAVE_PATH = MODEL_PATH / MODEL_NAME

# 3. Save the model state dict
torch.save(obj=model.state_dict(), f=MODEL_SAVE_PATH)

# Load a PyTorch Model

# Create a new instance
loaded_model = LinearRegressionModel()

# Load the saved model
loaded_model.load_state_dict(torch.load(MODEL_SAVE_PATH))

# Put the loaded model on a device
loaded_model.to(device)

# Evaluate loaded model
loaded_model.eval()
with torch.inference_mode():
loaded_model_preds = loaded_model(X_test.to(device))
loaded_model_preds == y_preds.to(device)

plot_predictions(X_train, y_train, X_test, y_test, predictions=loaded_model_preds.cpu())

You can access the code here on google collab.