Keras Tuner Subclass for Time Series Cross-Validation

Question

Custom Keras Tuner with Time Series Cross-Validation

I have written my own subclass of the default Keras tuner Tune class.

• Objective: I needed a way to incorporate time series cross-validation into the hyperparameter tuning process, which wasn't directly supported by the default Keras tuner.

• Functionality: My TimeSeriesBayesianOptimization subclass integrates time series cross-validation, allowing the model to be evaluated across multiple time-based splits and returning the average performance metrics.

• Use Case: This is particularly useful for my dataset, which involves time series forecasting where traditional random cross-validation can disrupt the temporal structure.

• Feedback Request: I'm looking for feedback on the efficiency of the implementation, potential pitfalls, and any best practices that I might have overlooked. I'm particularly interested in understanding if my method of averaging metrics across time series splits is optimal for guiding the Bayesian optimization process. Additionally, I want to ensure that the logic surrounding how the Oracle interprets the averaged objective over all the folds is sound.

The Tuner

class TimeSeriesBayesianOptimization(BayesianOptimization):
    def __init__(self, time_series_splits=5, *args, **kwargs):
        super(TimeSeriesBayesianOptimization, self).__init__(*args, **kwargs)
        self.time_series_splits = time_series_splits
        self.tscv = TimeSeriesSplit(n_splits=self.time_series_splits)

    def run_trial(self, trial, *args, **kwargs):
        # Extract X_train and y_train without altering original args
        X_train, y_train, *remaining_args = args

        # Callback to save the best epoch
        model_checkpoint = tuner_utils.SaveBestEpoch(
            objective=self.oracle.objective,
            filepath=self._get_checkpoint_fname(trial.trial_id),
        )
        original_callbacks = kwargs.pop("callbacks", [])

        # Track the histories
        histories = []

        for execution in range(self.executions_per_trial):
            total_val_loss = 0.0
            total_loss = 0.0
            total_binary_accuracy = 0.0
            total_val_binary_accuracy = 0.0

            for train_index, val_index in self.tscv.split(X_train):
                X_train_split, X_val_split = X_train[train_index], X_train[val_index]
                y_train_split, y_val_split = y_train[train_index], y_train[val_index]

                # Build the model for this trial's hyperparameters
                model = self.hypermodel.build(trial.hyperparameters)

                # Set up callbacks
                copied_callbacks = self._deepcopy_callbacks(original_callbacks)
                self._configure_tensorboard_dir(copied_callbacks, trial, execution)
                copied_callbacks.append(tuner_utils.TunerCallback(self, trial))
                copied_callbacks.append(model_checkpoint)

                # Train the model for this split
                history = model.fit(
                    X_train_split,
                    y_train_split,
                    validation_data=(X_val_split, y_val_split),
                    callbacks=copied_callbacks,
                    **kwargs
                )

                # Grab the best values for each metric
                best_val_loss = min(history.history["val_loss"])
                best_loss = min(history.history["loss"])
                best_binary_accuracy = max(history.history["binary_accuracy"])
                best_val_binary_accuracy = max(history.history["val_binary_accuracy"])

                # Accumulate the best values
                total_val_loss += best_val_loss
                total_loss += best_loss
                total_binary_accuracy += best_binary_accuracy
                total_val_binary_accuracy += best_val_binary_accuracy

            # Compute the averages
            avg_val_loss = total_val_loss / self.time_series_splits
            avg_loss = total_loss / self.time_series_splits
            avg_binary_accuracy = total_binary_accuracy / self.time_series_splits
            avg_val_binary_accuracy = (
                total_val_binary_accuracy / self.time_series_splits
            )

            # Store the averages in the histories list
            histories.append(
                {
                    "val_loss": avg_val_loss,
                    "loss": avg_loss,
                    "binary_accuracy": avg_binary_accuracy,
                    "val_binary_accuracy": avg_val_binary_accuracy,
                }
            )

        return histories

    def save_model(self, trial_id, model):
        """Save the model for the given trial."""
        fname = os.path.join(self.get_trial_dir(trial_id), "model.keras")
        model.save(fname)

Usage

# Cross Validation Search (WIP)
tuner = TimeSeriesBayesianOptimization(
    hypermodel=search_cnn_lstm_model,
    objective=Objective("val_loss", direction="min"),
    max_trials=6,
    time_series_splits=5,  # Number of time series splits for cross-validation
    seed=42,
    executions_per_trial=1,
    directory="tmp/tb",
    project_name="gru_cnn_vl",
)

tuner.search(
    X_train,
    y_train,
    epochs=120,
    shuffle=False,
    batch_size=72,
    callbacks=[
        tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=10, mode="min")
    ],
    verbose=True,
)