main.py

"""
Usage example:
    python main.py --data /path/to/dataset --architecture xresnet1d50 \
     --input-size 10 --fs-data 100 --fs-model 100
"""
import os

import torch
import lightning.pytorch as lp
from lightning.pytorch.tuner import Tuner
from lightning.pytorch.loggers import TensorBoardLogger
from lightning.pytorch.callbacks import (
    ModelCheckpoint,
    LearningRateMonitor,
    TQDMProgressBar,
    EarlyStopping,
)

from lite.cli import parse_args
from lite.movement import Main_Lite_Movement

# mlflow without autologging
# https://github.com/zjohn77/lightning-mlflow-hf/blob/74c30c784f719ea166941751bda24393946530b7/lightning_mlflow/train.py#L39
MLFLOW_AVAILABLE = True
try:
    import mlflow
    from lightning.pytorch.loggers import MLFlowLogger

    def log_params_from_namespace(hparams):
        # Handle both dict and Namespace objects
        params_dict = hparams if isinstance(hparams, dict) else hparams.__dict__
        for k, v in params_dict.items():
            mlflow.log_param(k, " " if str(v) == "" else str(v))

except ImportError:
    MLFLOW_AVAILABLE = False


def get_git_revision_short_hash():
    return ""  # subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD']).strip()


torch.set_float32_matmul_precision("high")


def main():
    hparams = parse_args()
    hparams.executable = "main_lite_" + hparams.modality
    hparams.revision = get_git_revision_short_hash()

    if not os.path.exists(hparams.output_path):
        os.makedirs(hparams.output_path)

    model = Main_Lite_Movement(hparams)

    logger = [TensorBoardLogger(save_dir=hparams.output_path, name="")]
    print("Output directory:", logger[0].log_dir)

    if MLFLOW_AVAILABLE and getattr(hparams, "mlflow", False):
        # Use custom experiment name if provided, otherwise use executable name
        experiment_name = (
            hparams.mlflow_experiment_name
            if getattr(hparams, "mlflow_experiment_name", None)
            else hparams.executable
        )
        mlflow.set_experiment(experiment_name)
        run = mlflow.start_run(run_name=hparams.metadata)
        mlf_logger = MLFlowLogger(
            experiment_name=mlflow.get_experiment(run.info.experiment_id).name,
            tracking_uri=mlflow.get_tracking_uri(),
            log_model=False,
        )
        mlf_logger._run_id = run.info.run_id
        mlf_logger.log_hyperparams = log_params_from_namespace
        logger.append(mlf_logger)

    checkpoint_callback = ModelCheckpoint(
        dirpath=logger[0].log_dir,
        filename="best_model",
        save_top_k=1,
        save_last=True,
        verbose=True,
        monitor=hparams.checkpoint_monitor,  # --checkpoint-monitor
        mode=hparams.checkpoint_mode,        # --checkpoint-mode
    )
    lr_monitor = LearningRateMonitor(logging_interval="step")
    callbacks = [checkpoint_callback, lr_monitor]

    if hparams.refresh_rate > 0:
        callbacks.append(TQDMProgressBar(refresh_rate=hparams.refresh_rate))

    # Add early stopping if requested
    if getattr(hparams, "early_stopping", False):
        early_stop_callback = EarlyStopping(
            monitor=hparams.early_stopping_monitor,
            min_delta=hparams.early_stopping_min_delta,
            patience=hparams.early_stopping_patience,
            mode=hparams.early_stopping_mode,
            verbose=True,
        )
        callbacks.append(early_stop_callback)
        print(f"Early stopping enabled: monitoring {hparams.early_stopping_monitor} "
              f"with patience={hparams.early_stopping_patience}, mode={hparams.early_stopping_mode}")

    trainer = lp.Trainer(
        num_sanity_val_steps=0,  # no debugging
        accumulate_grad_batches=hparams.accumulate,
        max_epochs=hparams.epochs,
        min_epochs=hparams.min_epochs,
        default_root_dir=hparams.output_path,
        logger=logger,
        callbacks=callbacks,
        benchmark=True,
        accelerator="gpu" if hparams.gpus > 0 else "cpu",
        devices=hparams.gpus if hparams.gpus > 0 else 1,
        num_nodes=hparams.num_nodes,
        precision=hparams.precision,
        enable_progress_bar=hparams.refresh_rate > 0,
    )

    if hparams.auto_batch_size:  # auto tune batch size
        Tuner(trainer).scale_batch_size(model, mode="binsearch")

    if hparams.lr_find:  # lr finder
        Tuner(trainer).lr_find(model)

    if hparams.epochs > 0 and hparams.eval_only == "":
        trainer.fit(model, ckpt_path=None if hparams.resume == "" else hparams.resume)
        # Preserve output_path from command-line args before testing
        model.hparams.output_path = hparams.output_path
        trainer.test(model, ckpt_path="best")
    elif hparams.eval_only != "":  # eval only
        # Preserve output_path from command-line args before testing
        model.hparams.output_path = hparams.output_path
        trainer.test(model, ckpt_path=hparams.eval_only)

    if MLFLOW_AVAILABLE and getattr(hparams, "mlflow", False):
        mlflow.end_run()


if __name__ == "__main__":
    main()