Skip to main content

Linear Regression with PyTorch

PyTorch is useful when linear regression is a building block in a larger differentiable model. For ordinary least squares with tabular data, scikit-learn is shorter and solves the optimization problem directly. Here, a single nn.Linear layer makes the equivalence explicit.

1. Prepare local data without leakage​

The synthetic fixture is deterministic. The scaler is fitted only on the training partition, then reused for the test partition.

import numpy as np
import torch
from sklearn.datasets import make_regression
from sklearn.metrics import root_mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

torch.manual_seed(42)
torch.set_num_threads(1)

X, y = make_regression(
    n_samples=240, n_features=4, noise=4.0, random_state=42
)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=42
)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)

X_train_t = torch.tensor(X_train, dtype=torch.float32)
y_train_t = torch.tensor(y_train[:, None], dtype=torch.float32)
X_test_t = torch.tensor(X_test, dtype=torch.float32)

model = torch.nn.Linear(X_train.shape[1], 1)
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.05)

model.train()
for _ in range(300):
    optimizer.zero_grad()
    loss = loss_fn(model(X_train_t), y_train_t)
    loss.backward()
    optimizer.step()

model.eval()
with torch.no_grad():
    y_pred = model(X_test_t).squeeze(1).numpy()

rmse = root_mean_squared_error(y_test, y_pred)
assert np.isfinite(rmse) and y_pred.shape == y_test.shape
print(f"Test RMSE: {rmse:.2f}")

2. Understand the API​

  • nn.Linear(p, 1) computes Xw+bXw+b for pp input features.
  • MSELoss optimizes mean squared error; it does not add regularization or uncertainty estimates.
  • zero_grad(), backward(), and step() form the training update.
  • model.eval() and torch.no_grad() are the correct inference context, even though this one-layer model has no dropout or batch normalization.

The fitted weights depend on feature scaling and optimization. If coefficient inference, exact least squares, or standard errors are the objective, use a statistical or classical linear-model implementation instead.

3. Production considerations​

Persist the training-fitted scaler with the model, validate input order and units, and monitor residuals after deployment. The fixed seed makes this tutorial reproducible on the pinned CPU stack; exact floating-point values can still vary across devices and library builds.