main.py

import os
import warnings  # Kept import
from datetime import datetime
import torch
import wandb
import hydra
from omegaconf import DictConfig, open_dict, omegaconf
from sklearn.model_selection import StratifiedShuffleSplit
import gc
import numpy as np

from utils.dataloader import init_dataloader
from utils.utils import init_setting, Logger
from training.training import build_model, Train

# --- START: Warning Suppression Block ---
# This is the correct way to suppress the warnings
os.environ["WANDB_SILENT"] = "true"
try:
    from urllib3.exceptions import NotOpenSSLWarning
    warnings.filterwarnings("ignore", category=NotOpenSSLWarning)
except ImportError:
    pass
# We set wandb=false in config, but this silences startup logs if it were on
# --- END: Warning Suppression Block ---

# Note: The 'wandb.require("core")' line was removed as it causes a warning.


@hydra.main(version_base=None, config_path="conf", config_name="config")
def main(cfg: DictConfig) -> None:
    # Initialize settings
    unique_id = datetime.now().strftime("%m-%d-%H-%M")
    init_setting(cfg)

    # -----------------------------------------------------------
    # START: DUMMY DATA GENERATION
    # Replace this section with your actual data loading when available.
    print("Generating dummy data for testing...")

    # Define the data dimensions based on the configuration
    num_train_valid_samples = 150
    num_test_samples = 50
    num_channels = cfg.dataset.n_channels
    num_times = cfg.dataset.n_times
    num_freq = cfg.dataset.n_frequencies
    num_classes = cfg.dataset.n_classes

    # Create dummy complex data
    # The model expects a tensor of shape [B, T, F, N, 2] (Batch, Time, Freq, Channels, 2)
    # The dataloader needs complex data and will create this shape
    train_valid_samples = np.random.rand(num_train_valid_samples, num_times, num_freq, num_channels) + \
                        1j * np.random.rand(num_train_valid_samples, num_times, num_freq, num_channels)
    train_valid_labs = np.random.randint(0, num_classes, num_train_valid_samples)
    train_valid_keys = train_valid_labs # Using labels as keys for StratifiedShuffleSplit

    test_samples = np.random.rand(num_test_samples, num_times, num_freq, num_channels) + \
                 1j * np.random.rand(num_test_samples, num_times, num_freq, num_channels)
    test_labs = np.random.randint(0, num_classes, num_test_samples)

    # END: DUMMY DATA GENERATION
    # -----------------------------------------------------------

    length = len(train_valid_samples) + len(test_samples)
    val_length = int(length * cfg.dataset.valid_set)
    with open_dict(cfg):
        cfg.val_length = val_length

    # setting record of experiments
    setting = 'B{}T{}S{}_E{}_T{}H{}F{}_S{}H{}F{}_E{}B{}_D{}_LR{}_D{}'.format(cfg.model.n_blocks,
                                                                             cfg.model.n_temporal_blocks,
                                                                             cfg.model.n_spatial_blocks,
                                                                             cfg.model.embed_dim,
                                                                             cfg.model.temporal_dmodel,
                                                                             cfg.model.temporal_nheads,
                                                                             cfg.model.temporal_dim_factor,
                                                                             cfg.model.spatial_dmodel,
                                                                             cfg.model.spatial_nheads,
                                                                             cfg.model.spatial_dim_factor,
                                                                             cfg.training.train_epochs,
                                                                             cfg.training.batch_size,
                                                                             cfg.model.dropout,
                                                                             cfg.optimizer.lr,
                                                                             cfg.optimizer.weight_decay
                                                                             )
    unique_id = '{}_{}'.format(unique_id, setting)
    path = os.path.join(cfg.checkpoints, unique_id)
    with open_dict(cfg):
        cfg.unique_id = unique_id
        cfg.path = path
    os.makedirs(path, exist_ok=True)

    logger = Logger(cfg)
    logger.init_logging()

    all_results = []
    # StratifiedShuffleSplit requires integer labels for classification
    split = StratifiedShuffleSplit(n_splits=cfg.n_fold, test_size=cfg.val_length, random_state=cfg.seed)
    for fold, (train_index, valid_index) in enumerate(split.split(train_valid_samples, train_valid_keys)):
        init_setting(cfg)
        if cfg.wandb:
            wandb.init(project=cfg.project,
                       group=unique_id,
                       name=f"{unique_id}:F{fold + 1}",
                       config=omegaconf.OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True),
                       reinit=True,
                       )
        print(f"<<<<<<<<<<<<< Fold[{fold + 1}/{cfg.n_fold}] >>>>>>>>>>>>>")
        dataloaders = init_dataloader(
            cfg=cfg,
            train_idx=train_index,
            valid_idx=valid_index,
            train_valid_samples=train_valid_samples,
            train_valid_labels=train_valid_labs,
            test_samples=test_samples,
            test_labels=test_labs
        )
        model = build_model(cfg)
        train_model = Train(
            cfg=cfg,
            model=model,
            dataloaders=dataloaders,
            fold=fold,
            logger=logger
        )
        final_result = train_model.train()
        all_results.append(final_result)
        
        if cfg.wandb:
            wandb.finish()

        gc.collect()
        # torch.empty_cache()

    logger.log_results(all_results=all_results)
    gc.collect()
    # torch.empty_cache()


if __name__ == '__main__':
    main()