saving trained models in optuna to a variable

Question

I have a project in which I do several optuna studies each having around 50 trials.

The optuna documentation suggests saving each model to a file for later use on this FAQ section

What I want is to have all the best models of different studies in a python list. How is that possible?

This is somewhat similar to my code:

def objective(trial, n_epochs, batch_size):
     params = {
              'learning_rate': trial.suggest_loguniform('learning_rate', 1e-5, 1e-1),
              'optimizer': trial.suggest_categorical("optimizer", ["Adam", "RMSprop", "SGD"]),
              'n_epochs': n_epochs,
              'batch_size': batch_size
              }
     model = clp_network(trial)
     accuracy, model = train_and_evaluate(params, model, trial) # THIS LINE
     return accuracy

     for i in range(50):
           study = optuna.create_study(direction="maximize", sampler=optuna.samplers.TPESampler(), pruner=optuna.pruners.MedianPruner())
           study.optimize(lambda trial: objective(trial, e, b), n_trials=50, show_progress_bar=True)

I would like to either save the model variable in the line marked THIS LINE, or somehow get the best model as a variable from the study.

Asked By: Mohsen Amiri

||

Source

Answer 1

The easiest way is to define a global variable to store a model for each trial as follows:

import optuna
from collections import defaultdict


models = defaultdict(dict)

def objective(t):
    model = t.suggest_int("x", 0, 100)
    models[t.study.study_name][t.number] = model
    
    return model

for _ in range(10):
    s = optuna.create_study()
    s.optimize(objective, n_trials=10)

However I reckon this approach is not scalable in terms of memory space, so I’d suggest removing non-best models after each optimize call or saving models on an external file as mentioned in Optuna’s FAQ.

Answered By: nzw0301

saving trained models in optuna to a variable

Question:

Answers: