pythonmachine-learningoptuna

saving trained models in optuna to a variable


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.


Solution

  • 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.