train_optuna.py

import os
import optuna
from data_gen import DataGenerator
from os import listdir
from utils import (
    iou,
    PlotLosses,
    dice_loss,
    focal_loss,
    categorical_loss,
    categorical_focal_loss,
    resolution2framesize3cha,
    resolution2framesize,
    bce_loss,
)
import matplotlib.pyplot as plt
import tensorflow as tf
from model.model import Thundernet as Thundernet_original
from models_repo.model_attention import Thundernet as Thundernet_attention
from models_repo.model_attention_2 import Thundernet as Thundernet_attention2
from models_repo.model_ppm_factors import Thundernet as Thundernet_ppm
from datetime import datetime
from matplotlib import pyplot as plt
from pathlib import Path

import os

# from data_gen_tfkeras import DataGenerator
from data_gen import DataGenerator
from os import listdir
from utils import (
    iou,
    PlotLosses,
    dice_loss,
    focal_loss,
    categorical_loss,
    categorical_focal_loss,
    resolution2framesize3cha,
    resolution2framesize,
)
import matplotlib.pyplot as plt
import tensorflow as tf

tf.config.run_functions_eagerly(True)
# from keras.backend.tensorflow_backend import set_session
import argparse
import sys
import numpy as np
import thundernet_config as Thundernet_config
from datetime import datetime
from matplotlib import pyplot as plt

from model.model import Thundernet as Thundernet_original
from models_repo.model_attention import Thundernet as Thundernet_attention
from models_repo.model_attention_2 import Thundernet as Thundernet_attention2
from models_repo.model_ppm_factors import Thundernet as Thundernet_ppm

from pathlib import Path
from collections import defaultdict
import copy

from collections import defaultdict

# Optuna-related imports
import optuna
import copy

plt.switch_backend("agg")


def objective(trial):
    # Define the hyperparameters you want to tune
    batch_size = trial.suggest_categorical("batch_size", [1, 2, 4])
    lr = trial.suggest_loguniform("lr", 1e-5, 1e-1)  # Learning rate
    kernel_regularizer = trial.suggest_loguniform("kernel_regularizer", 1e-5, 1e-2)

    # Call the main function with trial parameters
    return main(
        model="ppm",  # Use the 'ppm' model as per your request
        class_mappings=defaultdict(int, {1: 1}),
        batch_size=batch_size,
        lr=lr,
        kernel_regularizer=kernel_regularizer,
        epochs=1,  # Run only for 1 epoch
        loss="BCE",
        transformations=(),  # Add transformations as needed
    )


def main(
    model="original",
    class_mappings=None,
    batch_size=8,
    lr=1e-4,
    kernel_regularizer=0.001,
    epochs=1,
    loss="BCE",
    transformations=tuple(),
):

    # Parsing arguments for the main function
    FLAGS = argparse.Namespace(
        train_dir=Thundernet_config.train_path,
        val_dir=Thundernet_config.val_path,
        batch_size=batch_size,
        augment=Thundernet_config.augment,
        rand_crop=Thundernet_config.rand_crop,
        loss=loss,
        model_dir=Thundernet_config.model_dir,
        weights=Thundernet_config.weights,
        lr=lr,
        epochs=epochs,
        classes=Thundernet_config.classes,
        resolution=Thundernet_config.resolution,
        kernel_regularizer=kernel_regularizer,
        pretrained=Thundernet_config.pretrained_bool,
        pretrained_weigths=Thundernet_config.pretrained_weigths,
    )

    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"

    mypath_train = FLAGS.train_dir + "images/"
    label_path_train = FLAGS.train_dir + "labels/"
    list_IDs_train = [f[:-4] for f in listdir(mypath_train) if f[-4:] == ".jpg"]
    mypath_val = FLAGS.val_dir + "images/"
    label_path_val = FLAGS.val_dir + "labels/"
    list_IDs_val = [f[:-4] for f in listdir(mypath_val) if f[-4:] == ".jpg"]

    # Model Setup
    if model == "original":
        Thundernet = Thundernet_original
    elif model == "attention":
        Thundernet = Thundernet_attention
    elif model == "attention2":
        Thundernet = Thundernet_attention2
    elif model == "ppm":
        Thundernet = Thundernet_ppm
    else:
        raise ValueError(f"Unknown model: {model}")

    # Model directory setup
    model_dir = FLAGS.model_dir
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)

    thundernet = Thundernet(
        input_shape=resolution2framesize3cha(FLAGS.resolution),
        n_classes=FLAGS.classes,
        resnet_trainable=True,
        kernel_regularizer=FLAGS.kernel_regularizer,
    )

    if FLAGS.pretrained:
        thundernet.model.load_weights(
            FLAGS.pretrained_weigths, by_name=True, skip_mismatch=True
        )

    # Optimizer setup
    opt = tf.keras.optimizers.Adam(learning_rate=FLAGS.lr)

    # Set the loss function
    if FLAGS.loss == "BCE":
        loss = bce_loss()
    elif FLAGS.loss == "BFL":
        loss = focal_loss()
    elif FLAGS.loss == "DCL":
        loss = dice_loss()
    elif FLAGS.loss == "CFL":
        loss = categorical_focal_loss()
    elif FLAGS.loss == "CAT":
        loss = categorical_loss()

    thundernet.model.compile(loss=loss, optimizer=opt, metrics=[iou])

    # Data generators setup
    dataset_dir = Path(Thundernet_config.train_path).parent
    training_generator, validation_generator = DataGenerator.create_generators(
        dataset_dir,
        FLAGS.classes,
        training_batch_size=FLAGS.batch_size,
        to_stereo=False,
        transformations=transformations,
        class_mappings=class_mappings,
    )

    # Train the model
    history = thundernet.model.fit(
        training_generator,
        validation_data=validation_generator,
        epochs=FLAGS.epochs,
        class_weight=None,
        callbacks=[PlotLosses(model_dir)],
        use_multiprocessing=False,
        workers=6,
    )
    # Return validation loss or metric for Optuna optimization
    print(history)
    return np.mean(history.history["iou"])


# Optuna study setup
if __name__ == "__main__":
    study = optuna.create_study(
        direction="maximize", storage="sqlite:///db.sqlite3"
    )  # Minimize the validation loss
    study.optimize(objective, n_trials=100)  # Optimize for 10 trials
    print("Best hyperparameters found: ", study.best_params)

    import optuna.visualization as vis

    # Guardar el gráfico de importancia de parámetros
    fig = vis.plot_param_importances(study)
    fig.write_image("param_importance_IoU.png")

    # Guardar el gráfico del historial de optimización
    fig = vis.plot_optimization_history(study)
    fig.write_image("optimization_history_IoU.png")

    import pandas as pd

    # Assuming `study` is the Optuna study object
    df = study.trials_dataframe()

    df.to_csv("results_optuna_IoU.csv")

    # Plot Learning Rate vs Loss
    plt.figure(figsize=(8, 6))
    plt.scatter(df["params_lr"], df["value"], color="blue", alpha=0.7)
    plt.title("Learning Rate vs Loss")
    plt.xlabel("Learning Rate")
    plt.ylabel("Loss")
    plt.grid(True)
    plt.savefig("lr_vs_loss_IoU.png")
    plt.close()

    # Plot Weight Decay vs Loss
    plt.figure(figsize=(8, 6))
    plt.scatter(df["params_batch_size"], df["value"], color="green", alpha=0.7)
    plt.title("Batch size vs Loss")
    plt.xlabel("Batch size")
    plt.ylabel("Loss")
    plt.grid(True)
    plt.savefig("batch_size_vs_loss_IoU.png")
    plt.close()

    # Plot Loss Weight vs Loss
    plt.figure(figsize=(8, 6))
    plt.scatter(df["params_kernel_regularizer"], df["value"], color="red", alpha=0.7)
    plt.title("Kernel regularizer vs Loss")
    plt.xlabel("Kernel regularizer")
    plt.ylabel("Loss")
    plt.grid(True)
    plt.savefig("kernel_regularizer_vs_loss_IoU.png")
    plt.close()