Become a CreatorSign inGet Started

How to Use W&B Sweeps with LightGBM for Hyperparameter Tuning

In this article, I will walk through an example of how to use W&B Sweeps for hyperparameter tuning on LightGBM on the California Housing dataset available through scikit-learn.



3 months ago | 4 min read


Understand how different hyperparameters’ effects on model performance

Hyperparameter tuning is an important step in the modeling process to improve model performance and to customize model hyperparameters to better suit your dataset. There are different useful tools and packages that help with hyperparameter tuning, using grid, random or Bayesian searches.

These search functions return outputs of models trained on our dataset with different hyperparameters, but how do you make sense of the outputs from the hundreds or thousands of models trained on different permutations of hyperparameters?

Allow me to introduce Weights & Biases, a useful MLOps platform with tools to assist with model experiments tracking, model performance visualizations and more. The W&B Sweeps product is designed to help with hyperparameter tuning.

Some features to highlight include Bayesian optimization to prioritize which sets of hyperparameters to search through and powerful visualizations to help you understand the importance of different hyperparameters.

In this article, I will walk through an example of how to use W&B Sweeps for hyperparameter tuning on LightGBM on the California Housing dataset available through scikit-learn. The goal is to train a regression model to estimate value of houses in units of 100,000 in California given 8 different features.

I. Install W&B

Create an account on Weights & Biases and complete their simple sign up process.

Afterwards, install W&B in your Python environment, import it into your Python script or notebook and initialize the login procedure.

!pip install wandb

import wandb


You will then be prompted to enter your API key, which you can access via the W&B website after logging in under settings > API Keys.

II. Define Model Training Function

It is important to define a function that creates and trains the model to be used by the Sweep agent in later steps. This function will contain several components explained below:

  • Set Default configurations: specify default configurations in the form of a dictionary to start with that will be overwritten during the sweep.
  • Start W&B: Initialize a newwandb run.
  • Load and split the dataset: remember to add a seed to your train_test_split function for consistent splits across different runs.
  • Train Model: fit your model on the train set and use the trained model to predict on the test set.
  • Evaluate model performance: evaluate the predictions made by the trained model. In this case, I chose to use both mean_absolute_error and mean_squared_error as metrics for the regression model.
  • Log model performance metrics: record the metrics to W&B using wandb.log().

Below is the full Python code for this model training function:

III. Define Sweep Configurations

The next step is to define configurations for W&B Sweep to execute in the form of a dictionary. Some parameters to define here include:

  • method: specify your search strategy, examples include Bayesian, grid and random searches.
  • metric: define the name and goal (maximize or minimize) of the metric to use as the metric to optimize on. For example, I chose mean squared error (MSE) as my metric and the goal was to minimize this metric.
  • parameters: define the hyperparameters to tune stored as the keys of a dictionary and their corresponding values to search over in the form of a list stored as the values of this dictionary.

The resulting configurations of hyperparameters and their corresponding values to search over are defined below:

Configurations for W&B Sweep

See this page for a complete list of sweep configurations that you can customize.

IV. Initialize and Run Sweep with a W&B Agent

After defining the sweep configurations, we are ready to run the sweep. Begin by initializing the sweep to start a Sweep Controller. Pass in the dictionary of sweep configurations and a project name (string) as arguments in the initialized object as shown below:

sweep_id = wandb.sweep(sweep=sweep_configs, project="california-housing-sweeps")

Next, run the sweep agent and pass in both the sweep_id and model training function as arguments. You can also provide an optional argument to specify the total count of runs for the agent to make.

wandb.agent(sweep_id=sweep_id, function=train_model, count=30)

V. Visualize Results

Once your sweep has finished, you can go back to the W&B UI to look at the results. The two most useful visualizations for me are the hyperparameter importance plot and the parallel coordinates plot.

Hyperparameter Importance Plot

You can specify the metric to optimize on, and this chart will show you the rank-ordered results of which hyperparameters were most important. In our example dataset, the most important hyperparameter was by farn_estimators, and the negative correlation tells us that predictions by models with less trees have a higher MSE score. This makes sense because a model with less trees is likely underfitting our dataset.

Hyperparameter Importance Plot for LightGBM model trained on the California Housing dataset

Parallel Coordinates Plot

In Kaggle competitions, it is easy to pick the best set of hyperparameters returned from a search function, because you care most about optimizing your score and not so much about overfitting.

In reality, when models are used to make decisions that have an impact on a business’ bottom line, data scientists should and will care about optimizing the bias-variance tradeoff to avoid overfitting to the train set. The parallel coordinates plot is a very powerful visualization to help you understand the combined effects of different sets of hyperparameters.

As you can see from the plot below showing results from our example dataset, there are several different combinations of hyperparameters that can improve the model performance to the ~0.195 MSE score range, which is better than the baseline MSE score of 0.208 for a LightGBM model with default hyperparameters.

Parallel Coordinate Plot for LightGBM model trained on the California Housing dataset


W&B Sweeps is a powerful tool to assist data scientists with the hyperparameter tuning process.

The visualizations are useful in helping understand the effects of different sets of hyperparameters on the model and hence make more educated guesses to choose the set of hyperparameters to train a final model on. To see my python notebook, follow this link on Github.

I hope this was useful — please leave any comments below. Thanks!

Originally published here!


Created by



Data Science







Related Articles