Can’t make this part of the code run, and it’s really slow.
I just want to create a model to classify leafs images into 4 types (none, then 3 types of sickness)

I want to use F1 as a loss function, then use Bayesian optimisation to get the best parameters for my leaf class model but it’s not running. Or it’s going really slow then failing…
I’m running on cpu as i don’t have a GPU

# Objective function for Optuna
def objective(trial, train_dataloader, val_dataloader, device):
    # Define hyperparameter search space using Optuna's suggest functions
    conv1_filters = trial.suggest_categorical("conv1_filters", [16, 32, 64])
    conv2_filters = trial.suggest_categorical("conv2_filters", [64, 128, 256])
    kernel_size = trial.suggest_categorical("kernel_size", [3, 5, 7])
    hidden_units = trial.suggest_categorical("hidden_units", [256, 512, 1024])
    dropout_rate = trial.suggest_float("dropout_rate", 0.1, 0.5)
    learning_rate = trial.suggest_float("learning_rate", 1e-5, 1e-2, log=True)
    num_epochs = trial.suggest_int("num_epochs", 10, 50)

    # Initialize the model with the given params
    model = LeafCNN(
        conv1_filters=conv1_filters,
        conv2_filters=conv2_filters,
        kernel_size=kernel_size,
        hidden_units=hidden_units,
        dropout_rate=dropout_rate
    ).to(device)

    # Set up the optimizer and loss function
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    loss_fn = nn.CrossEntropyLoss(weight=class_weights_tensor)

    # Early stopping parameters
    patience = 10  # Allow 10 epochs without improvement
    delta = 0.001  # Minimum change that counts as an improvement
    best_val_loss = float("inf")
    patience_counter = 0  # Starting the counter at 0

    # Training loop
    for epoch in range(num_epochs):
        model.train()
        for X_batch, y_batch in train_dataloader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            optimizer.zero_grad()
            outputs = model(X_batch)
            loss = loss_fn(outputs, y_batch)
            loss.backward()
            optimizer.step()

        # Validation loss
        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for X_batch, y_batch in val_dataloader:
                X_batch, y_batch = X_batch.to(device), y_batch.to(device)
                outputs = model(X_batch)
                loss = loss_fn(outputs, y_batch)
                val_loss += loss.item()

        val_loss /= len(val_dataloader)
        print(f"Epoch {epoch + 1}/{num_epochs}, Val Loss: {val_loss}")

        if val_loss < best_val_loss - delta:
            best_val_loss = val_loss
            patience_counter = 0
        else:
            patience_counter += 1

        if patience_counter >= patience:
            print(f"Early stopping triggered at epoch {epoch + 1}")
            break

    # Evaluate using F1 score
    model.eval()
    all_preds, all_labels = [], []
    with torch.no_grad():
        for X_batch, y_batch in val_dataloader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            outputs = model(X_batch)
            _, predicted = torch.max(outputs, 1)
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(y_batch.cpu().numpy())

    f1 = f1_score(all_labels, all_preds, average="weighted")
    return f1  # We aim to maximize the F1 score

# Main function to run Optuna optimization
def optuna_search(train_dataloader, val_dataloader, device, num_trials):
    # Create Optuna study
    study = optuna.create_study(direction="maximize")  # Maximize the F1 score
    study.optimize(lambda trial: objective(trial, train_dataloader, val_dataloader, device), n_trials=num_trials)

    # Print the best hyperparameters
    print("Best hyperparameters found:", study.best_params)
    print("Best F1 Score:", study.best_value)

    return study.best_params, study.best_value


# Main Program
if __name__ == "__main__":

    # Create DataLoaders
    train_dataset = TensorDataset(X_train_split, y_train_split)
    train_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)

    val_dataset = TensorDataset(X_val, y_val)
    val_dataloader = DataLoader(val_dataset, batch_size=4, shuffle=False)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Run Optuna optimization
    best_params, best_f1 = optuna_search(train_dataloader, val_dataloader, device, num_trials=10)
    print("Best Params:", best_params)
    print("Best F1 Score:", best_f1)