In our first PyTorch model, we’ll implement many details ourselves as this will help us understand the model better.
For example, we’ll determine the predictions of the model by using matrix multiplication, implement the model parameters ourselves, and adjust the model parameters independently. The trained model parameters, slopes, and offsets denote the two most important learnable parameters within a neuron or a linear transformation.
Later, we’ll hand over these tasks more and more to the PyTorch framework. If we were to do this from the outset, many aspects of model training would remain black boxes that we wouldn’t fully understand.
Finally, we’ll train a model to predict the anxiety level (y) based on a variety of independent features, using the following formula:
y = w1 * X1 + w2 * X2 + … + w30 * X30 + b
We’ll start by importing and preparing the data. Then, we’ll train the model 2 and evaluate the training progress. Finally, we’ll check the model predictions.
We start as usual by importing the packages shown below. Since we’re building directly on the data preparation from a previous blog post, we import the independent features X and the dependent feature y directly from the Dataprep script. For the creation of tensors, we also load NumPy and torch, and for visualization, we load seaborn and matplotlib. Finally, we use the value to evaluate the model, and we therefore load the r2_score function from sklearn.
#%% packages
from DataPrep import X, y
import torch
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score
PyTorch only works with tensors, so we first convert the NumPy array into tensors with torch.from_numpy, as follows:
X_tensor = torch.from_numpy(X.astype(np.float32))
y_tensor = torch.from_numpy(y.astype(np.float32)) # Ensure y is float32
Now that we have the data in shape, we can get started. Our regression model is ultimately described by a bias parameter and a slope parameter (slope or weight). For each feature, there is a slope parameter and a total of one bias parameter.
We need to initialize the w (weight) and b (bias) terms first, and we can implement this with torch.zeros. We should also set the requires_grad parameter to True because that’s the only way to enable automatic backpropagation and training of the model.
# Initialize weights with smaller values to prevent exploding gradients
w = torch.zeros(X.shape[1], 1, requires_grad=True, dtype=torch.float32)
b = torch.zeros(1, requires_grad=True, dtype=torch.float32)
print(f"w shape: {w.shape}, b shape: {b.shape}")
w shape: torch.Size([30, 1]), b shape: torch.Size([1])
The training process is influenced by a number of parameters, and the most important ones are the number of epochs and the learning rate. Before we start the training, let’s take a closer look at these two important parameters.
Let’s start with the epochs. Our training dataset has 11,000 samples, and these are usually transferred to the model in smaller chunks called batches. We’ll come back to this concept later. Once all the samples have been used once to adjust the weights of the model, the epoch is complete, and the process is then repeated so that the model can “see” the same data many times to learn from it. Typically, a model is trained for several epochs and the patterns in the data are captured better with each successive epoch.
Now that we have an understanding of the concept of epochs, let’s clarify what the learning rate is all about. We can imagine model training as the search for the deepest point in an unknown valley. Say that a hiker descending from a mountaintop is blindfolded and must slowly feel his way down. They can decide whether their steps should be long or short, and the length of each step corresponds to the learning rate.
With long stride lengths (high learning rates), the hiker quickly covers a large area, but they could also quickly pass the lowest point and start climbing up the neighboring mountain. Conversely, they could choose to take very small steps (which correspond to low learning rates) in order to move very carefully. In that case, the hiker would have a high probability of finding the lowest point exactly and not going beyond it, but they could take a relatively long time to get there. However, they could also get stuck in a small, shallow puddle on the valley floor. The hiker could conclude that he has reached the deepest point, stop moving, and get stuck, even though the deepest point of the valley is still further away. The technical term for such an issue is a local minimum, which is in strong contrast to the deepest point (global minimum).
The learning rate therefore defines how quickly or carefully the blindfolded hiker explores the valley of the error to find the optimum point.
Now, we can really get started and define our two parameters:
EPOCHS = 100
LEARNING_RATE = 0.01
This brings us to the actual core of the model training—the training loop, which is implemented in the listing below. The data is shown to the model 100 times, and this is implemented with a for loop via the EPOCHS. We can then assess how well the model is learning by studying the losses, which are extracted in each epoch and added to the loss_list.
The loop always runs through the same steps, as follows:
loss_list = []
for epoch in range(EPOCHS):
# 1. Forward pass
y_predict = torch.matmul(X_tensor, w) + b
# 2. Calculate loss (MSE)
loss = torch.nn.functional.mse_loss(y_predict, y_tensor)
# 3. Backward pass
loss.backward()
# 4. Update weights and biases
with torch.no_grad():
w -= LEARNING_RATE * w.grad
b -= LEARNING_RATE * b.grad
# 5. Zero gradients after using them
w.grad.zero_()
b.grad.zero_()
# 6. Store loss for plotting
loss_list.append(loss.item())
# 7. Print loss for this epoch
print(f"Epoch {epoch}, Loss: {loss.item():.4f}")
Epoch 0, Loss: 19.9446
Epoch 1, Loss: 19.0938
...
Epoch 98, Loss: 1.5858
Epoch 99, Loss: 1.5732
This shows us how the training is progressing, and we can see that the losses are getting smaller.
Now, we can visualize the training losses again in a graphic. The corresponding code is shown in this listing, and the losses stored in the loss_list are shown as a line diagram above the number of EPOCHS. To do this, we use seaborn with the sns.lineplot function.
#%% plot loss
sns.lineplot(x=range(EPOCHS), y=loss_list)
plt.title('Loss over Epochs')
plt.xlabel('Epoch [-]')
plt.ylabel('Loss [-]')
The figure below shows the result of the model training. The loss decreases continuously with= each subsequent epoch, but you can also see here that although the model has fewer losses, the losses are asymptotically approaching a limit. We’ll return to the question of the optimum training duration at a later point.
The picture rarely looks as “clean” as it does here. There’s usually more fluctuation (i.e., periods when losses rise again slightly for a short time before returning to the longer-term falling trend).
Next, let’s take a closer look at the model weights (the slope values [w] and offset value [b]), as follows:
#%% check results
print(f"Weights: {w.detach().numpy().flatten()}, Bias: {b.item()}")
Weights: [-0.09916524 -0.5292779 -0.16298231 0.34272027 0.06862636
...
Bias: 3.408252716064453
Finally, in this simple case, we can use these values and perform the calculation based on the regression formula. We achieve this by multiplying the independent features by the model weights. We need to perform the calculation within the scope of torch.no_grad() to prevent gradients from being calculated. It’s also important to note that we’re performing model inference (i.e., testing the model) here, not model training. In model inference, we don’t want to perform any operations that could influence the network but should not be part of the training process. The positive side effects are that this saves resources (such as memory and computing time) and ensures that certain operations are not incorrectly included in the gradient calculation.
# %%
with torch.no_grad():
y_pred = (torch.matmul(X_tensor, w) + b).detach().numpy().flatten()
We’ve calculated the predictions y_pred and displayed them in connection with the real values y, and the listing below shows the corresponding code. We use the sns.regplot() function to create a scatterplot with a superimposed regression line. The data points are displayed in blue with a transparency value of 0.1, and that value ensures that the points are displayed in a bluer color in areas where many values lie on top of each other and the points in areas with very few points are displayed in a faint blue color. In addition, the regression line is displayed as a red line.
# %% visualize correlation
sns.regplot(x=y_pred, y=y, color='red',
scatter_kws={'s': 10,
'color': 'blue',
'alpha': 0.1})
plt.title('Predicted Anxiety Level vs Actual Anxiety Level')
plt.xlabel('Predicted Anxiety Level [-]')
plt.ylabel('Actual Anxiety Level [-]')
The result is a correlation diagram and can be seen in this figure. The actual anxiety level is plotted above the predicted anxiety level.
The correlation is positive, and on average, an anxiety level of 5 is also predicted as 5. But of course, there is scatter in the data, so in some cases, values between 1 and 7 are predicted. This illustration gives you a good overview of the areas in which the model works well and where improvements may still be necessary.
However, you’ll often want to compare different models with each other, and this is easier if you summarize the model quality as a single numerical value. In the field of regression models, the R2 value is a frequently used measure. The R2 value (also known as the coefficient of determination) is a statistical indicator that shows how well the independent features in a regression model match the variance (or dispersion) of the dependent variable. The value range is generally between 0 and 1 or 0% and 100%. We can understand the extreme values as follows:
If the R2 value is 0.75, for example, you can interpret it as meaning that 75% of the variance of the dependent variable can be explained by the independent variables contained in the model. The remaining 25% of the variance is due to other factors not included in the model or random errors.
In general, a higher R2 value reflects a better fit of the model to the data.
It’s also extremely important to note at this point that a high R2 value doesn’t necessarily mean that there is a causal relationship, meaning the independent variables do not automatically have a causal influence on the dependent variable. The high value merely shows a strong statistical correlation between the variables.
The calculation is performed using the r2_score function from sklearn, the real and predicted values are passed to the function, and a single numerical value is obtained, as follows:
r2 = r2_score(y_true=y,
y_pred=y_pred)
print(f"R-squared: {r2:.2f}")
R-squared: 0.65
Our first model achieves an R2 value of 0.65, and now, we can use this value as a benchmark to compare our model with other models.
Whether an R2 value is considered good or bad depends heavily on the context. In certain cases, an R2 of 0.98 is considered poor, and in other cases, an R2 of 0.4 is considered very good. But for now, we’re satisfied with the result, and we want to further improve our model training by using PyTorch’s capabilities to make our code more modular and therefore more versatile.
Editor’s note: This post has been adapted from a section of the book PyTorch: The Practical Guide by Bert Gollnick. Bert is a senior data scientist who specializes in renewable energies. For many years, he has taught courses about data science and machine learning, and more recently, about generative AI and natural language processing. Bert studied aeronautics at the Technical University of Berlin and economics at the University of Hagen. His main areas of interest are machine learning and data science.
This post was originally published 6/2026.