vasim.
Back to blog
← View series: ibm ai engineering

~/blog

Linear Regression One Parameter

Apr 1, 2026•4 min read•By Mohammed Vasim
AIMachine LearningLLMPyTorchTensorFlowGenerative AILangChainAI Agents

Skills Network LogoSkills Network Logo

Linear Regression 1D: Training One Parameter

Objective

  • How to create cost or criterion function using MSE (Mean Square Error).

Table of Contents

In this lab, you will train a model with PyTorch by using data that you created. The model only has one parameter: the slope.

Estimated Time Needed: 20 min

Preparation

The following are the libraries we are going to use for this lab.

python
# These are the libraries will be used for this lab.

import numpy as np
import matplotlib.pyplot as plt

The class plot_diagram helps us to visualize the data space and the parameter space during training and has nothing to do with PyTorch.

python
# The class for plotting

class plot_diagram():
    
    # Constructor
    def __init__(self, X, Y, w, stop, go = False):
        start = w.data
        self.error = []
        self.parameter = []
        print(type(X.numpy()))
        self.X = X.numpy()
       
        self.Y = Y.numpy()
        self.parameter_values = torch.arange(start, stop)
        self.Loss_function = [criterion(forward(X), Y) for w.data in self.parameter_values] 
        w.data = start
        
    # Executor
    def __call__(self, Yhat, w, error, n):
        self.error.append(error)
        self.parameter.append(w.data)
        plt.subplot(212)
        plt.plot(self.X, Yhat.detach().numpy())
        plt.plot(self.X, self.Y,'ro')
        plt.xlabel("A")
        plt.ylim(-20, 20)
        plt.subplot(211)
        plt.title("Data Space (top) Estimated Line (bottom) Iteration " + str(n))
        # Convert lists to PyTorch tensors
        parameter_values_tensor = torch.tensor(self.parameter_values)
        loss_function_tensor = torch.tensor(self.Loss_function)

        # Plot using the tensors
        plt.plot(parameter_values_tensor.numpy(), loss_function_tensor.numpy())
  
        plt.plot(self.parameter, self.error, 'ro')
        plt.xlabel("B")
        plt.figure()
    
    # Destructor
    def __del__(self):
        plt.close('all')

Make Some Data

Import PyTorch library:

python
# Import the library PyTorch

import torch

Generate values from -3 to 3 that create a line with a slope of -3. This is the line you will estimate.

python
# Create the f(X) with a slope of -3

X = torch.arange(-3, 3, 0.1).view(-1, 1)
f = -3 * X
python
X[:10], X.shape
python
f[:10], f.shape

Let us plot the line.

python
# Plot the line with blue

plt.plot(X.numpy(), f.numpy(), label = 'f')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.show()

Let us add some noise to the data in order to simulate the real data. Use torch.randn(X.size()) to generate Gaussian noise that is the same size as X and has a standard deviation opf 0.1.

python
# Add some noise to f(X) and save it in Y

Y = f + 0.1 * torch.randn(X.size())

Plot the Y:

python
# Plot the data points

plt.plot(X.numpy(), Y.numpy(), 'rx', label = 'Y')

plt.plot(X.numpy(), f.numpy(), label = 'f')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.show()

Create the Model and Cost Function (Total Loss)

In this section, let us create the model and the cost function (total loss) we are going to use to train the model and evaluate the result.

First, define the forward function $y=w*x$. (We will add the bias in the next lab.)

python
# Create forward function for prediction

def forward(x):
    return w * x

Define the cost or criterion function using MSE (Mean Square Error):

python
# Create the MSE function for evaluate the result.

def criterion(yhat, y):
    return torch.mean((yhat - y) ** 2)

Define the learning rate lr and an empty list LOSS to record the loss for each iteration:

python
# Create Learning Rate and an empty list to record the loss for each iteration

lr = 0.1
LOSS = []

Now, we create a model parameter by setting the argument requires_grad to True because the system must learn it.

python
w = torch.tensor(-10.0, requires_grad = True)

Create a plot_diagram object to visualize the data space and the parameter space for each iteration during training:

python
gradient_plot = plot_diagram(X, Y, w, stop = 5)

Train the Model

Let us define a function for training the model. The steps will be described in the comments.

python
# Define a function for train the model

def train_model(iter):
    for epoch in range (iter):
        
        # make the prediction as we learned in the last lab
        Yhat = forward(X)
        
        # calculate the iteration
        loss = criterion(Yhat,Y)
        
        # plot the diagram for us to have a better idea
        gradient_plot(Yhat, w, loss.item(), epoch)
        
        # store the loss into list
        LOSS.append(loss.item())
        
        # backward pass: compute gradient of the loss with respect to all the learnable parameters
        loss.backward()
        
        # updata parameters
        w.data = w.data - lr * w.grad.data
        
        # zero the gradients before running the backward pass
        w.grad.data.zero_()

Let us try to run 4 iterations of gradient descent:

python
# Give 4 iterations for training the model here.

train_model(4)

Plot the cost for each iteration:

python
# Plot the loss for each iteration

plt.plot(LOSS)
plt.tight_layout()
plt.xlabel("Epoch/Iterations")
plt.ylabel("Cost")

Practice

Create a new learnable parameter w with an initial value of -15.0.

python
# Practice: Create w with the inital value of -15.0

# Type your code here
w = torch.tensor(-15.0, requires_grad=True)

Double-click here for the solution.

Create an empty list LOSS2:

python
# Practice: Create LOSS2 list

# Type your code here
LOSS2 = []

Double-click here for the solution.

Write your own my_train_model function with loss list LOSS2. And run it with 4 iterations.

python
# Practice: Create your own my_train_model

gradient_plot2 = plot_diagram(X, Y, w, stop = 15)
python
def my_train_model(iter):

    for epoch in range(iter): 

        # Predict 
        yhat = forward(X)

        # Calculate the loss 
        loss = criterion(yhat, Y)

        # Plot training 
        gradient_plot2(yhat, w, loss.item(), epoch)

        # Store the loss in the list LOSS2 
        LOSS2.append(loss.item())

        # Backward 
        loss.backward()

        # Update params 
        w.data = w.data - lr * w.grad.data 

        # Zero gradient 
        w.grad.data.zero_()
python
my_train_model(4)

Double-click here for the solution.

Plot an overlay of the list LOSS2 and LOSS.

python
plt.plot(LOSS2)
plt.xlabel("Epoch")
plt.ylabel("Cost")
plt.show()
python
# Practice: Plot the list LOSS2 and LOSS

# Type your code here
plt.plot(LOSS2, label="LOSS")
plt.plot(LOSS, label="LOSS2")

plt.xlabel("Iteration/Epoch")
plt.ylabel("Cost")
plt.title("Loss Over Epochs")

plt.legend()

plt.show()

Double-click here for the solution.

What does this tell you about the parameter value?

Double-click here for the solution.

About the Authors:

Joseph Santarcangelo has a PhD in Electrical Engineering, his research focused on using machine learning, signal processing, and computer vision to determine how videos impact human cognition. Joseph has been working for IBM since he completed his PhD.

Other contributors: Michelle Carey, Mavis Zhou

© IBM Corporation. All rights reserved.

← PreviousPrediction 1D RegressionNext →Training Slope and Bias

Comments (0)

No comments yet. Be the first to comment!

Leave a comment