experiment.py

"""Exeriment"""

import argparse
import os
import signal
import sys

import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import (
    BatchSizeFinder,
    EarlyStopping,
    LearningRateMonitor,
    ModelCheckpoint,
)
from pytorch_lightning.loggers import TensorBoardLogger
from pytorch_lightning.plugins.environments import SLURMEnvironment

from src.har_project.data.build_dataset import DataModule
from src.har_project.models.models_factory import get_model
from src.har_project.trainer import TrainingModule

torch.set_float32_matmul_precision("medium")


def get_args():
    parser = argparse.ArgumentParser(
        "Fine-tuning and evaluation script for action classification",
        add_help=True,
    )

    # General Parameters
    parser.add_argument(
        "--batch_size",
        default=10,
        type=int,
        help="Batch size for training and evaluation.",
    )
    parser.add_argument(
        "--find_batch_size",
        action="store_true",
        default=False,
        help="Find the largest batch size for training.",
    )
    parser.add_argument(
        "--epochs", default=5, type=int, help="Number of epochs to train for."
    )
    parser.add_argument(
        "--save_ckpt_freq",
        default=1,
        type=int,
        help="Frequency (in epochs) to save checkpoints.",
    )
    parser.add_argument(
        "--save_top_k",
        default=1,
        type=int,
        help="Number of top checkpoints to save.",
    )
    parser.add_argument(
        "--output_dir",
        default="",
        type=str,
        help="Output directory for saving results. Defaults to project directory",
    )
    parser.add_argument(
        "--checkpoint_postfix",
        default="",
        type=str,
        help="Postfix for the checkpoint directory.",
    )

    # Model Parameters
    parser.add_argument(
        "--model_name",
        default="MoViNetA0",
        choices=[
            "MoViNetA0",
            # "MoViNetA0Stream",
            "x3d_xs",
            "x3d_s",
            "x3d_m",
            "cnn-rnn",
        ],
        type=str,
        help="Name of the model to use.",
    )
    parser.add_argument(
        "--input_size", default=172, type=int, help="video input size"
    )

    # Augmentation parameters
    parser.add_argument(
        "--aa",
        type=str,
        default="rand-m7-n4-mstd0.5-inc1",
        metavar="NAME",
        help='Use AutoAugment policy. "v0" or "original". " + "(default: rand-m7-n4-mstd0.5-inc1)',
    ),
    parser.add_argument(
        "--train_interpolation",
        type=str,
        default="bicubic",
        help='Training interpolation (random, bilinear, bicubic default: "bicubic")',
    )
    parser.add_argument(
        "--smoothing",
        type=float,
        default=0.1,
        help="Label smoothing (default: 0.1)",
    )
    parser.add_argument("--short_side_size", type=int, default=172)

    # Random Erase Params
    parser.add_argument(
        "--random_erasing_prob",
        type=float,
        default=0.0,
        help="Random erase prob (default: 0 (disabled))",
    )
    parser.add_argument(
        "--random_erasing_mode",
        type=str,
        default="pixel",
        help='Random erase mode (default: "pixel")',
    )
    parser.add_argument(
        "--random_erasing_max_count",
        type=int,
        default=1,
        help="Random erase count (default: 1)",
    )
    parser.add_argument(
        "--recount", type=int, default=1, help="Random erase count (default: 1)"
    )
    parser.add_argument(
        "--random_erasing_num_splits",
        action="store_true",
        default=False,
        help="Do not random erase first (clean) augmentation split",
    )

    # Optimizer Parameters
    parser.add_argument(
        "--lr",
        default=5e-5,
        type=float,
        help="Learning rate for the optimizer.",
    )
    parser.add_argument(
        "--warmup_lr",
        type=float,
        default=1e-6,
        metavar="LR",
        help="warmup learning rate (default: 1e-6)",
    )
    parser.add_argument(
        "--min_lr",
        type=float,
        default=1e-6,
        metavar="LR",
        help="lower lr bound for cyclic schedulers that hit 0 (1e-5)",
    )
    parser.add_argument(
        "--warmup_epochs",
        type=int,
        default=5,
        metavar="N",
        help="epochs to warmup LR, if scheduler supports",
    )
    parser.add_argument(
        "--weight_decay",
        type=float,
        default=0.05,
        help="weight decay (default: 0.05)",
    )
    parser.add_argument(
        "--opt_eps",
        default=1e-8,
        type=float,
        metavar="EPSILON",
        help="Optimizer Epsilon (default: 1e-8)",
    )

    # Dataset Parameters
    parser.add_argument(
        "--data_path", default="/path/to/dataset", type=str, help="dataset path"
    )
    parser.add_argument(
        "--data_set",
        default="UCF101",
        choices=[
            "UCF101",
            "HMDB51",
            "SSv2",
            "Kinetics400",
            "UCF101_smal",
            "UCF101_smal_frames",
            # "HMDB51",
        ],
        type=str,
        help="dataset",
    )
    parser.add_argument(
        "--nr_classes",
        default=101,
        type=int,
        help="number of the classification types",
    )
    parser.add_argument(
        "--fname_tmpl",
        default="{:05}.jpg",
        type=str,
        help="filename_tmpl for rawframe dataset",
    )
    parser.add_argument(
        "--start_idx",
        default=1,
        type=int,
        help="start index for rawframe dataset",
    )

    # Sampling Parameters
    parser.add_argument(
        "--sampling_strategy",
        type=str,
        default="uniform",
        choices=["dense", "uniform", "random"],
        help="Frame sampling strategy",
    )
    parser.add_argument(
        "--num_frames",
        type=int,
        default=50,
        help="Number of frames to sample per segment",
    )
    parser.add_argument(
        "--sampling_rate",
        type=int,
        default=4,
        help="How often we should sample frames for dense sampling",
    )
    parser.add_argument(
        "--num_sample", type=int, default=1, help="Repeated_aug (default: 1)"
    )
    parser.add_argument(
        "--keyframes",
        type=str,
        default="",
        help="Full path to the CSV file containing the keyframes.",
    )
    parser.add_argument(
        "--keyframes_prio",
        action="store_true",
        default=False,
        help="""Set this flag to consider keyframes in prioritized order.
        Defaults to False, treating all keyframes as equally important. 
        Will be ignored if --keyframes is not set.""",
    )
    parser.add_argument(
        "--eval_disable_keyframes",
        action="store_true",
        default=False,
        help="Disable keyframes during evaluation.",
    )

    # Training Options
    parser.add_argument(
        "--num_workers",
        default=6,
        type=int,
        help="Number of worker threads for DataLoader.",
    )
    parser.add_argument(
        "--pin_memory",
        default=True,
        type=bool,
        help="Whether to use pinned memory for DataLoader.",
    )
    parser.add_argument(
        "--checkpoint",
        default="",
        type=str,
        help="Path to the checkpoint for resuming training.",
    )
    parser.add_argument(
        "--eval",
        action="store_true",
        default=False,
        help="Evaluate the model.",
    )
    parser.add_argument(
        "--dev",
        action="store_true",
        default=False,
        help="Enable fast development run. If set, runs a single batch of train, validation, and test to check for any errors.",
    )
    parser.add_argument(
        "--limit_batches",
        action="store_true",
        default=False,
        help="Limit the number of batches for training, validation, and testing. Useful for quick iterations while degugging.",
    )
    parser.add_argument(
        "--seed",
        default=None,
        type=int,
        help="Seed value for random number generation.",
    )

    # Eval params
    parser.add_argument(
        "--extra_measurements",
        action="store_true",
        default=False,
        help="Enable measuring per video/frame latency (ms) and peak memory usage.",
    )

    # Slurm params
    parser.add_argument(
        "--nodes",
        default=1,
        type=int,
        help="Number of nodes to use for distributed training.",
    )
    parser.add_argument(
        "--gpus",
        default=1,
        type=int,
        help="Number of GPUs to use for training. Should match gres=gpu:n and ntasks-per-node=n",
    )
    parser.add_argument(
        "--srequeue",
        action="store_true",
        default=False,
        help="Enable automatic requeueing in SLURM environment.",
    )

    return parser.parse_args(args=None if sys.argv[1:] else ["--help"])


def main(args):
    """Main"""
    # Seed everything
    pl.seed_everything(args.seed)

    # build checkpoint directory name
    name = f"{args.data_set}_{args.model_name}"
    if args.checkpoint_postfix != "":
        name += "_" + args.checkpoint_postfix

    project_dir_path = (
        os.path.dirname(os.path.abspath(__file__))
        if args.output_dir == ""
        else args.output_dir
    )
    checkpoint_dir = os.path.join(project_dir_path, "checkpoints", name)

    if args.extra_measurements:
        print(
            "===> Measuring extrac metrics: latency and memory usage. Setting batch size to 1 and forcing eval mode."
        )
        args.batch_size = 1
        args.eval = True

    # import dataset
    data_module = DataModule(
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        pin_memory=args.pin_memory,
        args=args,
    )

    # Import backbone
    backbone = get_model(args.model_name, args.nr_classes)

    # define the model
    model = TrainingModule(
        model=backbone,
        num_classes=args.nr_classes,
        label_smoothing=args.smoothing,
        lr=args.lr,
        lr_min=args.min_lr,
        eps=args.opt_eps,
        weight_decay=args.weight_decay,
        warmup_epochs=args.warmup_epochs,
        warmup_lr=args.warmup_lr,
    )

    # define callbacks
    callbacks = [
        ModelCheckpoint(
            dirpath=checkpoint_dir,
            filename="{epoch}_{val_loss:.3f}_{val_accTop1:.3f}",
            save_top_k=args.save_top_k,
            monitor="val_loss",
            mode="min",
            save_last=True,
        ),
        EarlyStopping(
            monitor="val_loss",
            min_delta=1e-4,
            patience=10,
            verbose=False,
            mode="min",
        ),
        LearningRateMonitor(logging_interval="step"),
    ]
    if args.find_batch_size:
        callbacks.append(
            BatchSizeFinder(mode="binsearch", init_val=args.batch_size)
        )

    # Train
    trainer_params = {
        "accelerator": "auto",
        "precision": 16,
        "max_epochs": args.epochs,
        "log_every_n_steps": 1,
        "fast_dev_run": args.dev,
        "callbacks": callbacks,
        "logger": TensorBoardLogger(
            save_dir=checkpoint_dir, name="lightning_logs"
        ),
    }
    if args.limit_batches:
        trainer_params.update(
            {
                "limit_train_batches": 0.1,
                "limit_val_batches": 0.2,
                "limit_test_batches": 0.3,
            }
        )
    if args.nodes > 1 or args.gpus > 1:
        trainer_params.update(
            {
                "strategy": "ddp",
                "devices": args.gpus,
                "num_nodes": args.nodes,
            }
        )

    # slurm parameters
    slurm_env = SLURMEnvironment(
        auto_requeue=args.srequeue, requeue_signal=signal.SIGHUP
    )
    if slurm_env.detect():
        slurm_env.validate_settings(args.gpus, args.nodes)

    trainer = pl.Trainer(
        **trainer_params,
        plugins=[slurm_env],
    )

    # Start training or set checkpoint path for evaluation
    best_model_path = None
    if args.eval:
        if not args.checkpoint:
            raise ValueError(
                "The --eval flag requires a --checkpoint to be specified."
            )
        best_model_path = args.checkpoint
    else:
        trainer.fit(
            model=model,
            datamodule=data_module,
            ckpt_path=(
                args.checkpoint if os.path.isfile(args.checkpoint) else None
            ),  # If args.checkpoint is provided, it will resume training
        )

        # Get path to best model
        best_model_path = trainer.checkpoint_callback.best_model_path

    # Evaluate on test set
    if not args.dev:  # disable evaluation during development
        if not os.path.isfile(best_model_path):
            raise ValueError(f"Checkpoint not found at {best_model_path}.")

        # Load best checkpoint
        print(f"Loading checkpoint: {best_model_path}")
        model = TrainingModule.load_from_checkpoint(
            checkpoint_path=best_model_path,
            model=backbone,
            extra_measurements=args.extra_measurements,
        )
        trainer.test(model=model, dataloaders=data_module)


if __name__ == "__main__":
    args = get_args()
    main(args)