Gabriele Cosmo
217 lines
11 KiB
Python
217 lines
11 KiB
Python
from dataclasses import dataclass
|
|
from typing import Tuple, Dict, Any, List
|
|
|
|
import numpy as np
|
|
from optuna import Trial, create_study, get_all_study_summaries, load_study
|
|
from optuna.pruners import MedianPruner
|
|
from optuna.samplers import TPESampler
|
|
from optuna.trial import TrialState
|
|
|
|
from core.constants import LEARNING_RATE, BATCH_SIZE_PER_REPLICA, ACTIVATION, OUT_ACTIVATION, \
|
|
OPTIMIZER_TYPE, KERNEL_INITIALIZER, BIAS_INITIALIZER, N_TRIALS, LATENT_DIM, \
|
|
INTERMEDIATE_DIMS, MAX_HIDDEN_LAYER_DIM, GLOBAL_CHECKPOINT_DIR
|
|
from core.model import VAEHandler
|
|
from utils.preprocess import preprocess
|
|
|
|
|
|
@dataclass
|
|
class HyperparameterTuner:
|
|
"""Tuner which looks for the best hyperparameters of a Variational Autoencoder specified in model.py.
|
|
|
|
Currently, supported hyperparameters are: dimension of latent space, number of hidden layers, learning rate,
|
|
activation function, activation function after the final layer, optimizer type, kernel initializer,
|
|
bias initializer, batch size.
|
|
|
|
Attributes:
|
|
_discrete_parameters: A dictionary of hyperparameters taking discrete values in the range [low, high].
|
|
_continuous_parameters: A dictionary of hyperparameters taking continuous values in the range [low, high].
|
|
_categorical_parameters: A dictionary of hyperparameters taking values specified by the list of them.
|
|
_storage: A string representing URL to a database required for a distributed training
|
|
_study_name: A string, a name of study.
|
|
|
|
"""
|
|
_discrete_parameters: Dict[str, Tuple[int, int]]
|
|
_continuous_parameters: Dict[str, Tuple[float, float]]
|
|
_categorical_parameters: Dict[str, List[Any]]
|
|
_storage: str = None
|
|
_study_name: str = None
|
|
|
|
def _check_hyperparameters(self):
|
|
available_hyperparameters = ["latent_dim", "nb_hidden_layers", "learning_rate", "activation", "out_activation",
|
|
"optimizer_type", "kernel_initializer", "bias_initializer",
|
|
"batch_size_per_replica"]
|
|
hyperparameters_to_be_optimized = list(self._discrete_parameters.keys()) + list(
|
|
self._continuous_parameters.keys()) + list(self._categorical_parameters.keys())
|
|
for hyperparameter_name in hyperparameters_to_be_optimized:
|
|
if hyperparameter_name not in available_hyperparameters:
|
|
raise Exception(f"Unknown hyperparameter: {hyperparameter_name}")
|
|
|
|
def __post_init__(self):
|
|
self._check_hyperparameters()
|
|
self._energies_train, self._cond_e_train, self._cond_angle_train, self._cond_geo_train = preprocess()
|
|
|
|
if self._storage is not None and self._study_name is not None:
|
|
# Parallel optimization
|
|
study_summaries = get_all_study_summaries(self._storage)
|
|
if any(self._study_name == study_summary.study_name for study_summary in study_summaries):
|
|
# The study is already created in the database. Load it.
|
|
self._study = load_study(self._study_name, self._storage)
|
|
else:
|
|
# The study does not exist in the database. Create a new one.
|
|
self._study = create_study(storage=self._storage, sampler=TPESampler(), pruner=MedianPruner(),
|
|
study_name=self._study_name, direction="minimize")
|
|
else:
|
|
# Single optimization
|
|
self._study = create_study(sampler=TPESampler(), pruner=MedianPruner(), direction="minimize")
|
|
|
|
def _create_model_handler(self, trial: Trial) -> VAEHandler:
|
|
"""For a given trail builds the model.
|
|
|
|
Optuna suggests parameters like dimensions of particular layers of the model, learning rate, optimizer, etc.
|
|
|
|
Args:
|
|
trial: Optuna's trial
|
|
|
|
Returns:
|
|
Variational Autoencoder (VAE)
|
|
"""
|
|
|
|
# Discrete parameters
|
|
if "latent_dim" in self._discrete_parameters.keys():
|
|
latent_dim = trial.suggest_int(name="latent_dim",
|
|
low=self._discrete_parameters["latent_dim"][0],
|
|
high=self._discrete_parameters["latent_dim"][1])
|
|
else:
|
|
latent_dim = LATENT_DIM
|
|
|
|
if "nb_hidden_layers" in self._discrete_parameters.keys():
|
|
nb_hidden_layers = trial.suggest_int(name="nb_hidden_layers",
|
|
low=self._discrete_parameters["nb_hidden_layers"][0],
|
|
high=self._discrete_parameters["nb_hidden_layers"][1])
|
|
|
|
all_possible = np.arange(start=latent_dim + 5, stop=MAX_HIDDEN_LAYER_DIM)
|
|
chunks = np.array_split(all_possible, nb_hidden_layers)
|
|
ranges = [(chunk[0], chunk[-1]) for chunk in chunks]
|
|
ranges = reversed(ranges)
|
|
|
|
# Cast from np.int to int allows to become JSON serializable.
|
|
intermediate_dims = [trial.suggest_int(name=f"intermediate_dim_{i}", low=int(low), high=int(high)) for
|
|
i, (low, high)
|
|
in enumerate(ranges)]
|
|
else:
|
|
intermediate_dims = INTERMEDIATE_DIMS
|
|
|
|
if "batch_size_per_replica" in self._discrete_parameters.keys():
|
|
batch_size_per_replica = trial.suggest_int(name="batch_size_per_replica",
|
|
low=self._discrete_parameters["batch_size_per_replica"][0],
|
|
high=self._discrete_parameters["batch_size_per_replica"][1])
|
|
else:
|
|
batch_size_per_replica = BATCH_SIZE_PER_REPLICA
|
|
|
|
# Continuous parameters
|
|
if "learning_rate" in self._continuous_parameters.keys():
|
|
learning_rate = trial.suggest_float(name="learning_rate",
|
|
low=self._continuous_parameters["learning_rate"][0],
|
|
high=self._continuous_parameters["learning_rate"][1])
|
|
else:
|
|
learning_rate = LEARNING_RATE
|
|
|
|
# Categorical parameters
|
|
if "activation" in self._categorical_parameters.keys():
|
|
activation = trial.suggest_categorical(name="activation",
|
|
choices=self._categorical_parameters["activation"])
|
|
else:
|
|
activation = ACTIVATION
|
|
|
|
if "out_activation" in self._categorical_parameters.keys():
|
|
out_activation = trial.suggest_categorical(name="out_activation",
|
|
choices=self._categorical_parameters["out_activation"])
|
|
else:
|
|
out_activation = OUT_ACTIVATION
|
|
|
|
if "optimizer_type" in self._categorical_parameters.keys():
|
|
optimizer_type = trial.suggest_categorical(name="optimizer_type",
|
|
choices=self._categorical_parameters["optimizer_type"])
|
|
else:
|
|
optimizer_type = OPTIMIZER_TYPE
|
|
|
|
if "kernel_initializer" in self._categorical_parameters.keys():
|
|
kernel_initializer = trial.suggest_categorical(name="kernel_initializer",
|
|
choices=self._categorical_parameters["kernel_initializer"])
|
|
else:
|
|
kernel_initializer = KERNEL_INITIALIZER
|
|
|
|
if "bias_initializer" in self._categorical_parameters.keys():
|
|
bias_initializer = trial.suggest_categorical(name="bias_initializer",
|
|
choices=self._categorical_parameters["bias_initializer"])
|
|
else:
|
|
bias_initializer = BIAS_INITIALIZER
|
|
|
|
checkpoint_dir = f"{GLOBAL_CHECKPOINT_DIR}/{self._study_name}/trial_{trial.number:03d}"
|
|
|
|
return VAEHandler(_wandb_project_name=self._study_name,
|
|
_wandb_tags=["hyperparameter tuning", f"trial {trial.number}"],
|
|
_batch_size_per_replica=batch_size_per_replica,
|
|
_intermediate_dims=intermediate_dims,
|
|
latent_dim=latent_dim,
|
|
_learning_rate=learning_rate,
|
|
_activation=activation,
|
|
_out_activation=out_activation,
|
|
_optimizer_type=optimizer_type,
|
|
_kernel_initializer=kernel_initializer,
|
|
_bias_initializer=bias_initializer,
|
|
_checkpoint_dir=checkpoint_dir,
|
|
_early_stop=True,
|
|
_save_model_every_epoch=False,
|
|
_save_best_model=True,
|
|
)
|
|
|
|
def _objective(self, trial: Trial) -> float:
|
|
"""For a given trial trains the model and returns an average validation loss.
|
|
|
|
Args:
|
|
trial: Optuna's trial
|
|
|
|
Returns: One float numer which is a validation loss. It can be either calculated as an average of k trainings
|
|
performed in cross validation mode or is one number obtained from validation on unseen before, some fraction
|
|
of the dataset.
|
|
"""
|
|
|
|
# Generate the trial model.
|
|
model_handler = self._create_model_handler(trial)
|
|
|
|
# Train the model.
|
|
verbose = True
|
|
histories = model_handler.train(self._energies_train, self._cond_e_train, self._cond_angle_train,
|
|
self._cond_geo_train, verbose)
|
|
|
|
# Return validation loss (currently it is treated as an objective goal). Notice that we take into account the
|
|
# best model according to the validation loss.
|
|
final_validation_losses = [np.min(history.history["val_loss"]) for history in histories]
|
|
avg_validation_loss = np.mean(final_validation_losses).item()
|
|
return avg_validation_loss
|
|
|
|
def tune(self) -> None:
|
|
"""Main tuning function.
|
|
|
|
Based on a given study, tunes the model and prints detailed information about the best trial (value of the
|
|
objective function and adjusted parameters).
|
|
"""
|
|
|
|
self._study.optimize(func=self._objective, n_trials=N_TRIALS, gc_after_trial=True)
|
|
pruned_trials = self._study.get_trials(deepcopy=False, states=(TrialState.PRUNED,))
|
|
complete_trials = self._study.get_trials(deepcopy=False, states=(TrialState.COMPLETE,))
|
|
print("Study statistics: ")
|
|
print(" Number of finished trials: ", len(self._study.trials))
|
|
print(" Number of pruned trials: ", len(pruned_trials))
|
|
print(" Number of complete trials: ", len(complete_trials))
|
|
|
|
print("Best trial:")
|
|
trial = self._study.best_trial
|
|
|
|
print(" Value: ", trial.value)
|
|
|
|
print(" Params: ")
|
|
for key, value in trial.params.items():
|
|
print(f" {key}: {value}")
|