Python How to retrieve the best model from Optuna LightGBM study?

Python

Optuna

LightGBM

Model Selection

Machine Learning

Python How to retrieve the best model from Optuna LightGBM study?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Machine Learning Python

Introduction

Optuna optimizes hyperparameters by running multiple trials, each training a model with different settings. To retrieve the best model from an Optuna-LightGBM study, you access study.best_trial.params and retrain the model with those parameters, or you store the trained model (via callbacks or artifact logging) during the objective function so you can load it directly without retraining. The key insight is that Optuna tracks parameter values and scores but does not store the trained model object by default — you must handle model persistence yourself.

Basic Optuna + LightGBM Setup

python

1import optuna
2import lightgbm as lgb
3from sklearn.datasets import load_breast_cancer
4from sklearn.model_selection import train_test_split
5
6data = load_breast_cancer()
7X_train, X_test, y_train, y_test = train_test_split(
8    data.data, data.target, test_size=0.2, random_state=42
9)
10
11def objective(trial):
12    params = {
13        'objective': 'binary',
14        'metric': 'binary_logloss',
15        'verbosity': -1,
16        'boosting_type': 'gbdt',
17        'num_leaves': trial.suggest_int('num_leaves', 20, 300),
18        'learning_rate': trial.suggest_float('learning_rate', 1e-3, 0.3, log=True),
19        'feature_fraction': trial.suggest_float('feature_fraction', 0.4, 1.0),
20        'bagging_fraction': trial.suggest_float('bagging_fraction', 0.4, 1.0),
21        'min_child_samples': trial.suggest_int('min_child_samples', 5, 100),
22    }
23
24    dtrain = lgb.Dataset(X_train, label=y_train)
25    cv_results = lgb.cv(params, dtrain, nfold=5, num_boost_round=200,
26                        callbacks=[lgb.early_stopping(20)])
27    best_score = cv_results['valid binary_logloss-mean'][-1]
28    return best_score
29
30study = optuna.create_study(direction='minimize')
31study.optimize(objective, n_trials=100)

Retrieving Best Parameters and Retraining

The most common approach — get the best parameters and retrain:

python

1# Get the best trial's parameters
2best_params = study.best_trial.params
3print(f"Best score: {study.best_value:.4f}")
4print(f"Best params: {best_params}")
5
6# Retrain with the best parameters
7final_params = {
8    'objective': 'binary',
9    'metric': 'binary_logloss',
10    'verbosity': -1,
11    'boosting_type': 'gbdt',
12    **best_params  # Merge the optimized parameters
13}
14
15dtrain = lgb.Dataset(X_train, label=y_train)
16dtest = lgb.Dataset(X_test, label=y_test, reference=dtrain)
17
18best_model = lgb.train(
19    final_params,
20    dtrain,
21    num_boost_round=200,
22    valid_sets=[dtest],
23    callbacks=[lgb.early_stopping(20)]
24)
25
26# Use the model
27predictions = best_model.predict(X_test)

Storing the Model During Optimization

To avoid retraining, save the model inside the objective function:

python

1import joblib
2import os
3
4os.makedirs('models', exist_ok=True)
5
6def objective_with_save(trial):
7    params = {
8        'objective': 'binary',
9        'metric': 'binary_logloss',
10        'verbosity': -1,
11        'num_leaves': trial.suggest_int('num_leaves', 20, 300),
12        'learning_rate': trial.suggest_float('learning_rate', 1e-3, 0.3, log=True),
13        'feature_fraction': trial.suggest_float('feature_fraction', 0.4, 1.0),
14    }
15
16    dtrain = lgb.Dataset(X_train, label=y_train)
17    cv_results = lgb.cv(params, dtrain, nfold=5, num_boost_round=200,
18                        callbacks=[lgb.early_stopping(20)])
19    best_score = cv_results['valid binary_logloss-mean'][-1]
20    best_iteration = len(cv_results['valid binary_logloss-mean'])
21
22    # Train on full training set and save
23    model = lgb.train(params, dtrain, num_boost_round=best_iteration)
24    model.save_model(f'models/trial_{trial.number}.txt')
25
26    return best_score
27
28study = optuna.create_study(direction='minimize')
29study.optimize(objective_with_save, n_trials=50)
30
31# Load the best model directly
32best_model = lgb.Booster(model_file=f'models/trial_{study.best_trial.number}.txt')
33predictions = best_model.predict(X_test)

Using Optuna's LightGBM Tuner

Optuna provides a specialized LightGBMTunerCV that handles the tuning automatically:

python

1import optuna.integration.lightgbm as lgb_tuner
2
3params = {
4    'objective': 'binary',
5    'metric': 'binary_logloss',
6    'verbosity': -1,
7    'boosting_type': 'gbdt',
8}
9
10dtrain = lgb.Dataset(X_train, label=y_train)
11
12tuner = lgb_tuner.LightGBMTunerCV(
13    params, dtrain, nfold=5, num_boost_round=200,
14    callbacks=[lgb.early_stopping(20)],
15    optuna_seed=42
16)
17tuner.run()
18
19print(f"Best score: {tuner.best_score}")
20print(f"Best params: {tuner.best_params}")
21
22# Retrain with best params from tuner
23best_model = lgb.train(tuner.best_params, dtrain, num_boost_round=200,
24                       callbacks=[lgb.early_stopping(20)],
25                       valid_sets=[lgb.Dataset(X_test, label=y_test)])

Using Trial User Attributes

Store arbitrary metadata (including model references) on each trial:

python

1def objective_with_attrs(trial):
2    params = {
3        'objective': 'binary',
4        'metric': 'binary_logloss',
5        'verbosity': -1,
6        'num_leaves': trial.suggest_int('num_leaves', 20, 300),
7        'learning_rate': trial.suggest_float('learning_rate', 1e-3, 0.3, log=True),
8    }
9
10    dtrain = lgb.Dataset(X_train, label=y_train)
11    cv_results = lgb.cv(params, dtrain, nfold=5, num_boost_round=200,
12                        callbacks=[lgb.early_stopping(20)])
13
14    best_iteration = len(cv_results['valid binary_logloss-mean'])
15    trial.set_user_attr('best_iteration', best_iteration)
16
17    return cv_results['valid binary_logloss-mean'][-1]
18
19study = optuna.create_study(direction='minimize')
20study.optimize(objective_with_attrs, n_trials=50)
21
22# Retrieve stored metadata
23best_iteration = study.best_trial.user_attrs['best_iteration']
24best_model = lgb.train({**study.best_trial.params, 'objective': 'binary'},
25                       lgb.Dataset(X_train, label=y_train),
26                       num_boost_round=best_iteration)

Common Pitfalls

Assuming study.best_trial returns a trained model: Optuna stores parameters and scores, not model objects. You must retrain with study.best_trial.params or save the model inside the objective function and load it afterward.
Forgetting fixed parameters when retraining: study.best_trial.params only contains the parameters Optuna tuned (via trial.suggest_*). Fixed parameters like objective and metric must be merged back manually with {**fixed_params, **study.best_trial.params}.
Not matching num_boost_round during retraining: The optimal number of boosting rounds from early stopping in CV is lost unless you store it via trial.set_user_attr(). Retraining with a different round count gives different results.
Using study.best_params vs study.best_trial.params: Both return the same parameters, but study.best_trial also gives access to user_attrs, number, value, and datetime_start — all useful for logging and model management.
Disk space from saving every trial's model: Saving a model file per trial adds up quickly with hundreds of trials. Either save only when a new best score is found (if trial.number == study.best_trial.number) or clean up non-best models after the study completes.

Summary

Access optimized parameters via study.best_trial.params and retrain the model with those values
Save models inside the objective function to avoid retraining — use model.save_model() or joblib.dump()
Use trial.set_user_attr() to store metadata like best_iteration for faithful reproduction
optuna.integration.lightgbm.LightGBMTunerCV provides a streamlined interface for LightGBM-specific tuning
Always merge fixed parameters (objective, metric) with Optuna's tuned parameters when retraining

Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course

Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Python How to retrieve the best model from Optuna LightGBM study?

ML System Design practice on Codemia

Introduction

Basic Optuna + LightGBM Setup

Retrieving Best Parameters and Retraining

Storing the Model During Optimization

Using Optuna's LightGBM Tuner

Using Trial User Attributes

Common Pitfalls

Summary

Related reading

Tackling System Design Interview Problems

Track what you have practised

ML System Design practice on Codemia