Building a Regression Model in PyTorch

PyTorch library is for deep learning. Some applications of deep learning models are to solve regression or classification problems.
In this post, you will discover how to use PyTorch to develop and evaluate neural network models for regression problems.

After completing this post, you will know:

  • How to load data from scikit-learn and adapt it for PyTorch models
  • How to create a neural network for regerssion problem using PyTorch
  • How to improve model performance with data preparation techniques

Kick-start your project with my book Deep Learning with PyTorch. It provides self-study tutorials with working code.

Let’s get started.

Building a Regression Model in PyTorchPhoto by Sam Deng. Some rights reserved.

Description of the Dataset

The dataset you will use in this tutorial is the California housing dataset.

This is a dataset that describes the median house value for California districts. Each data sample is a census block group. The target variable is the median house value in USD 100,000 in 1990 and there are 8 input features, each describing something about the house. They are, namely,

  • MedInc: median income in block group
  • HouseAge: median house age in block group
  • AveRooms: average number of rooms per household
  • AveBedrms: average number of bedrooms per household
  • Population: block group population
  • AveOccup: average number of household members
  • Latitude: block group centroid latitude
  • Longitude: block group centroid longitude

This data is special because the input data is in vastly different scale. For example, the number of rooms per house is usually small but the population per block group is usually large. Moreover, most features should be positive but the longitude must be negative (because that’s about California). Handling such diversity of data is a challenge to some machine learning models.

You can get the dataset from scikit-learn, which in turn, is downloaded from the Internet at realtime:

Building a Model and Train

This is a regression problem. Unlike classification problems, the output variable is a continuous value. In case of neural networks, you usually use linear activation at the output layer (i.e., no activation) such that the output range theoretically can be anything from negative infinty to positive infinity.

Also for regression problems, you should never expect the model to predict the values perfectly. Therefore, you should care about how close the prediction is to the actual value. The loss metric that you can use for this is the mean square error (MSE) or mean absolute error (MAE). But you may also interested in the root mean squared error (RMSE) because that’s a metric in the same unit as your output variable.

Let’s try the traditional design of a neural network, namely, the pyramid structure. A pyramid structure is to have the number of neurons in each layer decreasing as the network progresses to the output. The number of input features is fixed, but you set a large number of neurons on the first hidden layer and gradually reduce the number in the subsequent layers. Because you have only one target in this dataset, the final layer should output only one value.

One design is as follows:

To train this network, you need to define a loss function. MSE is a reasonable choice. You also need an optimizer, such as Adam.

To train this model, you can use your usual training loop. In order to obtain an evaluation score so you are confident that the model works, you need to split the data into training and test sets. You may also want to avoid overfitting by keeping track on the test set MSE. The following is the training loop with the train-test split:

In the training loop, tqdm is used to set up a progress bar and in each iteration step, MSE is calculated and reported. You can see how the MSE changed by setting the tqdm parameter disable above to False.

Note that in the training loop, each epoch is to run the forward and backward steps with the training set a few times to optimize the model weights, and at the end of the epoch, the model is evaluated using the test set. It is the MSE from the test set that is remembered in the list history. It is also the metric to evaluate a model, which the best one is stored in the variable best_weights.

After you run this, you will have the best model restored and the best MSE stored in the variable best_mse. Note that the mean square error is the average of the square of the difference between the predicted value and the actual value. The square root of it, RMSE, can be regarded as the average difference and it is numerically more useful.

In below, you can show the MSE and RMSE, and plot the history of MSE. It should be decreasing with the epochs.

This model produced:

The MSE graph would like the following.

Putting everything together, the following is the complete code.

Want to Get Started With Deep Learning with PyTorch?

Take my free email crash course now (with sample code).

Click to sign-up and also get a free PDF Ebook version of the course.

Improving the Model with Preprocessing

In the above, you see the RMSE is 0.68. Indeed, it is easy to improve the RMSE by polishing the data before training. The problem of this dataset is the diversity of the features: Some are with a narrow range and some are wide. And some are small but positive while some are very negative. This indeed is not very nice to most of the machine learning model.

One way to improve this is to apply a standard scaler. It is to convert each feature into their standard score. In other words, for each feature $x$, you replace it with

z = \frac{x – \bar{x}}{\sigma_x}

Where $\bar{x}$ is the mean of $x$ and $\sigma_x$ is the standard deviation. This way, every transformed feature is centered around 0 and in a narrow range that around 70% of the samples are between -1 to +1. This can help the machine learning model to converge.

You can apply the standard scaler from scikit-learn. The following is how you should modify the data preparation part of the above code:

Note that standard scaler is applied after train-test split. The StandardScaler above is fitted on the training set but applied on both the training and test set. You must not apply the standard scaler to all data because nothing from the test set should be hinted to the model. Otherwise you are introducing data leakage.

Other than that, virtually nothing shall be changed: You still have 8 features (only they are not the same in value). You still use the same training loop. If you train the model with the scaled data, you should see the RMSE improved, e.g.,

While the MSE history is in a similar falling shape, the y-axis shows it is indeed better after scaling:

However, you need to be careful at the end: When you use the trained model and apply to new data, you should apply the scaler to the input data before feed into the mode. That is, inference should be done as follows:

The following is the complete code:

Of course, there is still room to imporve the model. One way is to present the target in log scale or, equivalently, use mean absolute percentage error (MAPE) as the loss function. This is because the target variable is the value of houses and it is in a wide range. For the same error magnitude, it is more an issue for low-valued houses. It is your exercise to modify the above code to produce a better prediction.


In this post, you discovered the use of PyTorch to build a regression model.

You learned how you can work through a regression problem step-by-step with PyTorch, specifically:

  • How to load and prepare data for use in PyTorch
  • How to create neural network models and choose a loss function for regression
  • How to improve model accuracy by applying standard scaler

Get Started on Deep Learning with PyTorch!

Deep Learning with PyTorch

Learn how to build deep learning models

...using the newly released PyTorch 2.0 library

Discover how in my new Ebook:
Deep Learning with PyTorch

It provides self-study tutorials with hundreds of working code to turn you from a novice to expert. It equips you with
tensor operation, training, evaluation, hyperparameter optimization, and much more...

Kick-start your deep learning journey with hands-on exercises

See What's Inside

9 Responses to Building a Regression Model in PyTorch

  1. Avatar
    George February 20, 2023 at 1:13 am #

    Re Model Evaluation

    Lines 92 and 935 respectfully are:
    with torch.no_grad():
    # Test out inference with 5 samples
    for i in range(5):

    And the validation evaluation excludes the code “with torch.no_grad():”
    see line 76:

    Why is used for one section for the evaluation and not the other?
    The pytorch discussion website states that the code “with torch.no_grad():” can be used for both.

    • Avatar
      James Carmichael February 20, 2023 at 10:06 am #

      Hi George…No particular reason or limitation. You should make the portions consistent and let us know what you find.

  2. Avatar
    George February 23, 2023 at 2:12 am #

    Hi James,

    It’s not how the program is run or when the evaluation is called. That won’t tell me a lot of information.

    What you need to know is what the pie torch documentation says rather than rely on how a particular model runs.

  3. Avatar
    Arnold February 25, 2023 at 3:33 am #

    Is there a way to turn off the intermediate output while iterating through the epochs? Thanks.

  4. Avatar
    Dimple May 26, 2023 at 5:32 pm #

    Hello James,

    I am writing the Regression Model in PyTorch for multiple inputs and multiple outputs, 1st I did the normalization (transformation) of inputs and outputs, but in the de-normalization stage, my cade is showing the error, can you please help me to debug this?
    TypeError: ‘int’ object is not callable

  5. Avatar
    Jorge Osorio June 1, 2023 at 5:43 am #

    thank you very much for all your tutorials, they helped me a lot, you are great.

    • Avatar
      James Carmichael June 2, 2023 at 10:04 am #

      You are very welcome Jorge! We greatly appreciate your support and feedback!

Leave a Reply