configs back at target level

configs/callbacks/early_stopping.yaml

@@ -1,6 +1,7 @@
 # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.EarlyStopping.html
 
 early_stopping:
+  _target_: pytorch_lightning.callbacks.EarlyStopping
   monitor: val_loss # quantity to be monitored, must be specified !!!
   min_delta: 0. # minimum change in the monitored quantity to qualify as an improvement
   patience: 3 # number of checks with no improvement after which training will be stopped

configs/callbacks/model_checkpoint.yaml

@@ -1,6 +1,7 @@
 # https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html
 
 model_checkpoint:
+  _target_: pytorch_lightning.callbacks.ModelCheckpoint
   dirpath: null # directory to save the model file
   filename: best-checkpoint # checkpoint filename
   monitor: val_loss # name of the logged metric which determines when model is improving

configs/callbacks/rich_model_summary.yaml

@@ -1,2 +1,3 @@
 rich_model_summary:
+  _target_: pytorch_lightning.callbacks.RichModelSummary
   max_depth: 1

configs/callbacks/rich_progress_bar.yaml

@@ -1,2 +1,3 @@
 rich_progress_bar:
+  _target_: pytorch_lightning.callbacks.RichProgressBar
   refresh_rate: 1

configs/data/catdog.yaml

@@ -1,5 +1,4 @@
 _target_: src.datamodules.catdog_datamodule.CatDogImageDataModule
-
 data_dir: ${paths.data_dir}
 url: ${paths.data_url}
 num_workers: 4

configs/logger/aim.yaml

@@ -1,4 +1,5 @@
 aim:
+  _target_: aim.pytorch_lightning.AimLogger
   experiment: ${name}
   train_metric_prefix: train_
   test_metric_prefix: test_

configs/logger/csv.yaml

@@ -1,6 +1,7 @@
 # csv logger built in lightning
 
 csv:
+  _target_: lightning.pytorch.loggers.csv_logs.CSVLogger
   save_dir: "${paths.output_dir}"
   name: "csv/"
   prefix: ""

configs/logger/mlflow.yaml

@@ -1,6 +1,7 @@
 # MLflow logger configuration
 
 mlflow:
+  _target_: lightning.pytorch.loggers.MLFlowLogger
   experiment_name: ${name}
   tracking_uri: file:${paths.log_dir}/mlruns
   save_dir: ${paths.log_dir}/mlruns

configs/logger/tensorboard.yaml

@@ -1,9 +1,10 @@
 # https://www.tensorflow.org/tensorboard/
 
 tensorboard:
+  _target_: lightning.pytorch.loggers.tensorboard.TensorBoardLogger
   save_dir: "${paths.output_dir}/tensorboard/"
   name: null
   log_graph: False
   default_hp_metric: True
   prefix: ""
-  # version: ""
+  # version: ""

notebooks/datamodule_lightning.ipynb

@@ -0,0 +1,92 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "vscode": {
+     "languageId": "plaintext"
+    }
+   },
+   "source": [
+    "In this notebook, we will be discussing about the pytorch lightning datamodule library with images in a folder strutcture with folders as class labels. We will be using  the cats and dogs dataset from kaggle. The dataset can be downloaded from [here](https://www.kaggle.com/c/dogs-vs-cats/data). The dataset contains 25000 images of cats and dogs. We will be using 20000 images for training and 5000 images for validation. The images are in a folder structure with folders as class labels."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/javascript": "IPython.notebook.set_autosave_interval(300000)"
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Autosaving every 300 seconds\n"
+     ]
+    }
+   ],
+   "source": [
+    "%autosave 300\n",
+    "%load_ext autoreload\n",
+    "%autoreload 2\n",
+    "%reload_ext autoreload\n",
+    "%config Completer.use_jedi = False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/mnt/batch/tasks/shared/LS_root/mounts/clusters/soutrik-vm-dev/code/Users/Soutrik.Chowdhury/pytorch-template-aws\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "import os\n",
+    "\n",
+    "os.chdir(\"..\")\n",
+    "print(os.getcwd())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "emlo_env",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.15"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

src/train.py

@@ -1,3 +1,9 @@
+"""
+Train and evaluate a model using PyTorch Lightning.
+Initializes the DataModule, Model, Trainer, and runs training and testing.
+Initializes loggers and callbacks from the configuration using Hydra configuration but with a more modular approach without direct instantiation.
+"""
+
 import os
 import shutil
 from pathlib import Path

src/train_new.py

@@ -0,0 +1,198 @@
+"""
+Train and evaluate a model using PyTorch Lightning.
+Initializes the DataModule, Model, Trainer, and runs training and testing.
+Initializes loggers and callbacks from the configuration using Hydra and target paths from the configuration.
+"""
+
+import os
+import shutil
+from pathlib import Path
+from typing import List
+import torch
+import lightning as L
+from dotenv import load_dotenv, find_dotenv
+import hydra
+from omegaconf import DictConfig, OmegaConf
+from src.utils.logging_utils import setup_logger, task_wrapper
+from loguru import logger
+import rootutils
+from lightning.pytorch.loggers import Logger
+from lightning.pytorch.callbacks import Callback
+
+# Load environment variables
+load_dotenv(find_dotenv(".env"))
+
+# Setup root directory
+
+root = rootutils.setup_root(__file__, indicator=".project-root")
+
+
+def instantiate_callbacks(callback_cfg: DictConfig) -> List[Callback]:
+    """Instantiate and return a list of callbacks from the configuration."""
+    callbacks_ls: List[L.Callback] = []
+
+    if not callback_cfg:
+        logger.warning("No callback configs found! Skipping..")
+        return None
+
+    if not isinstance(callback_cfg, DictConfig):
+        raise TypeError("Callbacks config must be a DictConfig!")
+
+    for _, cb_conf in callback_cfg.items():
+        if "_target_" in cb_conf:
+            logger.info(f"Instantiating callback <{cb_conf._target_}>")
+            callbacks_ls.append(hydra.utils.instantiate(cb_conf))
+
+    return callbacks_ls
+
+
+def instantiate_loggers(logger_cfg: DictConfig) -> List[Logger]:
+    """Instantiate and return a list of loggers from the configuration."""
+    loggers_ls: List[Logger] = []
+
+    if not logger_cfg:
+        logger.warning("No logger configs found! Skipping..")
+        return loggers_ls
+
+    if not isinstance(logger_cfg, DictConfig):
+        raise TypeError("Logger config must be a DictConfig!")
+
+    for _, lg_conf in logger_cfg.items():
+        if "_target_" in lg_conf:
+            logger.info(f"Instantiating logger <{lg_conf._target_}>")
+            loggers_ls.append(hydra.utils.instantiate(lg_conf))
+
+    return loggers_ls
+
+
+def load_checkpoint_if_available(ckpt_path: str) -> str:
+    """Return the checkpoint path if available, else None."""
+    if ckpt_path and Path(ckpt_path).exists():
+        logger.info(f"Using checkpoint: {ckpt_path}")
+        return ckpt_path
+    logger.warning(f"Checkpoint not found at {ckpt_path}. Using current model weights.")
+    return None
+
+
+def clear_checkpoint_directory(ckpt_dir: str):
+    """Clear checkpoint directory contents without removing the directory."""
+    ckpt_dir_path = Path(ckpt_dir)
+    if not ckpt_dir_path.exists():
+        logger.info(f"Creating checkpoint directory: {ckpt_dir}")
+        ckpt_dir_path.mkdir(parents=True, exist_ok=True)
+    else:
+        logger.info(f"Clearing checkpoint directory: {ckpt_dir}")
+        for item in ckpt_dir_path.iterdir():
+            try:
+                item.unlink() if item.is_file() else shutil.rmtree(item)
+            except Exception as e:
+                logger.error(f"Failed to delete {item}: {e}")
+
+
+@task_wrapper
+def train_module(
+    data_module: L.LightningDataModule, model: L.LightningModule, trainer: L.Trainer
+):
+    """Train the model and log metrics."""
+    logger.info("Starting training")
+    trainer.fit(model, data_module)
+    train_metrics = trainer.callback_metrics
+    train_acc = train_metrics.get("train_acc")
+    val_acc = train_metrics.get("val_acc")
+    logger.info(
+        f"Training completed. Metrics - train_acc: {train_acc}, val_acc: {val_acc}"
+    )
+    return train_metrics
+
+
+@task_wrapper
+def run_test_module(
+    cfg: DictConfig,
+    datamodule: L.LightningDataModule,
+    model: L.LightningModule,
+    trainer: L.Trainer,
+):
+    """Test the model using the best checkpoint or current model weights."""
+    logger.info("Starting testing")
+    datamodule.setup(stage="test")
+    test_metrics = trainer.test(
+        model, datamodule, ckpt_path=load_checkpoint_if_available(cfg.ckpt_path)
+    )
+    logger.info(f"Test metrics: {test_metrics}")
+    return test_metrics[0] if test_metrics else {}
+
+
+@hydra.main(config_path="../configs", config_name="train", version_base="1.1")
+def setup_run_trainer(cfg: DictConfig):
+    """Set up and run the Trainer for training and testing."""
+    # Display configuration
+    logger.info(f"Config:\n{OmegaConf.to_yaml(cfg)}")
+
+    # Initialize logger
+    log_path = Path(cfg.paths.log_dir) / (
+        "train.log" if cfg.task_name == "train" else "eval.log"
+    )
+    setup_logger(log_path)
+
+    # Display key paths
+    for path_name in [
+        "root_dir",
+        "data_dir",
+        "log_dir",
+        "ckpt_dir",
+        "artifact_dir",
+        "output_dir",
+    ]:
+        logger.info(
+            f"{path_name.replace('_', ' ').capitalize()}: {cfg.paths[path_name]}"
+        )
+
+    # Initialize DataModule and Model
+    logger.info(f"Instantiating datamodule <{cfg.data._target_}>")
+    datamodule: L.LightningDataModule = hydra.utils.instantiate(cfg.data)
+    logger.info(f"Instantiating model <{cfg.model._target_}>")
+    model: L.LightningModule = hydra.utils.instantiate(cfg.model)
+
+    # Check GPU availability and set seed for reproducibility
+    logger.info("GPU available" if torch.cuda.is_available() else "No GPU available")
+    L.seed_everything(cfg.seed, workers=True)
+
+    # Set up callbacks, loggers, and Trainer
+    callbacks = instantiate_callbacks(cfg.callbacks)
+    logger.info(f"Callbacks: {callbacks}")
+    loggers = instantiate_loggers(cfg.loggers)
+    logger.info(f"Loggers: {loggers}")
+    trainer: L.Trainer = hydra.utils.instantiate(
+        cfg.trainer, callbacks=callbacks, logger=loggers
+    )
+
+    # Training phase
+    train_metrics = {}
+    if cfg.get("train"):
+        clear_checkpoint_directory(cfg.paths.ckpt_dir)
+        train_metrics = train_module(datamodule, model, trainer)
+        (Path(cfg.paths.ckpt_dir) / "train_done.flag").write_text(
+            "Training completed.\n"
+        )
+
+    # Testing phase
+    test_metrics = {}
+    if cfg.get("test"):
+        test_metrics = run_test_module(cfg, datamodule, model, trainer)
+
+    # Combine metrics and extract optimization metric
+    all_metrics = {**train_metrics, **test_metrics}
+    optimization_metric = all_metrics.get(cfg.get("optimization_metric"), 0.0)
+    (
+        logger.warning(
+            f"Optimization metric '{cfg.get('optimization_metric')}' not found. Defaulting to 0."
+        )
+        if optimization_metric == 0.0
+        else logger.info(f"Optimization metric: {optimization_metric}")
+    )
+
+    return optimization_metric
+
+
+if __name__ == "__main__":
+    setup_run_trainer()