datamodule new tested

configs/data/catdog.yaml

@@ -1,8 +1,9 @@
 _target_: src.datamodules.catdog_datamodule.CatDogImageDataModule
-data_dir: ${paths.data_dir}
+root_dir: ${paths.data_dir}
+data_dir: "cats_and_dogs_filtered"
 url: ${paths.data_url}
 num_workers: 4
 batch_size: 32
 train_val_split: [0.8, 0.2]
 pin_memory: False
-image_size: 160
+image_size: 224

configs/experiment/catdog_experiment.yaml

@@ -18,10 +18,11 @@ seed: 42
 name: "catdog_experiment"
 
 data:
-  batch_size: 64
+  dataset: "cats_and_dogs_filtered"
+  batch_size: 32
   num_workers: 8
   pin_memory: True
-  image_size: 160
+  image_size: 224
 
 model:
   lr: 1e-3

notebooks/datamodule_lightning.ipynb

@@ -53,13 +53,222 @@
     }
    ],
    "source": [
-    "\n",
     "import os\n",
     "\n",
     "os.chdir(\"..\")\n",
     "print(os.getcwd())"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/anaconda/envs/emlo_env/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "from pathlib import Path\n",
+    "from typing import Union, Tuple, Optional, List\n",
+    "import os\n",
+    "import lightning as L\n",
+    "from torch.utils.data import DataLoader, random_split\n",
+    "from torchvision import transforms\n",
+    "from torchvision.datasets import ImageFolder\n",
+    "from torchvision.datasets.utils import download_and_extract_archive\n",
+    "from loguru import logger"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class CatDogImageDataModule(L.LightningDataModule):\n",
+    "    \"\"\"DataModule for Cat and Dog Image Classification using ImageFolder.\"\"\"\n",
+    "\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        data_root: Union[str, Path] = \"data\",\n",
+    "        data_dir: Union[str, Path] = \"cats_and_dogs_filtered\",\n",
+    "        batch_size: int = 32,\n",
+    "        num_workers: int = 4,\n",
+    "        train_val_split: List[float] = [0.8, 0.2],\n",
+    "        pin_memory: bool = False,\n",
+    "        image_size: int = 224,\n",
+    "        url: str = \"https://download.pytorch.org/tutorials/cats_and_dogs_filtered.zip\",\n",
+    "    ):\n",
+    "        super().__init__()\n",
+    "        self.data_root = Path(data_root)\n",
+    "        self.data_dir = data_dir\n",
+    "        self.batch_size = batch_size\n",
+    "        self.num_workers = num_workers\n",
+    "        self.train_val_split = train_val_split\n",
+    "        self.pin_memory = pin_memory\n",
+    "        self.image_size = image_size\n",
+    "        self.url = url\n",
+    "\n",
+    "        # Initialize variables for datasets\n",
+    "        self.train_dataset = None\n",
+    "        self.val_dataset = None\n",
+    "        self.test_dataset = None\n",
+    "\n",
+    "    def prepare_data(self):\n",
+    "        \"\"\"Download the dataset if it doesn't exist.\"\"\"\n",
+    "        self.dataset_path = self.data_root / self.data_dir\n",
+    "        if not self.dataset_path.exists():\n",
+    "            logger.info(\"Downloading and extracting dataset.\")\n",
+    "            download_and_extract_archive(\n",
+    "                url=self.url, download_root=self.data_root, remove_finished=True\n",
+    "            )\n",
+    "            logger.info(\"Download completed.\")\n",
+    "\n",
+    "    def setup(self, stage: Optional[str] = None):\n",
+    "        \"\"\"Set up the train, validation, and test datasets.\"\"\"\n",
+    "\n",
+    "        train_transform = transforms.Compose(\n",
+    "            [\n",
+    "                transforms.Resize((self.image_size, self.image_size)),\n",
+    "                transforms.RandomHorizontalFlip(0.1),\n",
+    "                transforms.RandomRotation(10),\n",
+    "                transforms.RandomAffine(0, shear=10, scale=(0.8, 1.2)),\n",
+    "                transforms.RandomAutocontrast(0.1),\n",
+    "                transforms.RandomAdjustSharpness(2, 0.1),\n",
+    "                transforms.ToTensor(),\n",
+    "                transforms.Normalize(\n",
+    "                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]\n",
+    "                ),\n",
+    "            ]\n",
+    "        )\n",
+    "\n",
+    "        test_transform = transforms.Compose(\n",
+    "            [\n",
+    "                transforms.Resize((self.image_size, self.image_size)),\n",
+    "                transforms.ToTensor(),\n",
+    "                transforms.Normalize(\n",
+    "                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]\n",
+    "                ),\n",
+    "            ]\n",
+    "        )\n",
+    "\n",
+    "        train_path = self.dataset_path / \"train\"\n",
+    "        test_path = self.dataset_path / \"test\"\n",
+    "\n",
+    "        self.prepare_data()\n",
+    "\n",
+    "        if stage == \"fit\" or stage is None:\n",
+    "            full_train_dataset = ImageFolder(root=train_path, transform=train_transform)\n",
+    "            self.class_names = full_train_dataset.classes\n",
+    "            train_size = int(self.train_val_split[0] * len(full_train_dataset))\n",
+    "            val_size = len(full_train_dataset) - train_size\n",
+    "            self.train_dataset, self.val_dataset = random_split(\n",
+    "                full_train_dataset, [train_size, val_size]\n",
+    "            )\n",
+    "            logger.info(\n",
+    "                f\"Train/Validation split: {len(self.train_dataset)} train, {len(self.val_dataset)} validation images.\"\n",
+    "            )\n",
+    "\n",
+    "        if stage == \"test\" or stage is None:\n",
+    "            self.test_dataset = ImageFolder(root=test_path, transform=test_transform)\n",
+    "            logger.info(f\"Test dataset size: {len(self.test_dataset)} images.\")\n",
+    "\n",
+    "    def _create_dataloader(self, dataset, shuffle: bool = False) -> DataLoader:\n",
+    "        \"\"\"Helper function to create a DataLoader.\"\"\"\n",
+    "        return DataLoader(\n",
+    "            dataset=dataset,\n",
+    "            batch_size=self.batch_size,\n",
+    "            num_workers=self.num_workers,\n",
+    "            pin_memory=self.pin_memory,\n",
+    "            shuffle=shuffle,\n",
+    "        )\n",
+    "\n",
+    "    def train_dataloader(self) -> DataLoader:\n",
+    "        return self._create_dataloader(self.train_dataset, shuffle=True)\n",
+    "\n",
+    "    def val_dataloader(self) -> DataLoader:\n",
+    "        return self._create_dataloader(self.val_dataset)\n",
+    "\n",
+    "    def test_dataloader(self) -> DataLoader:\n",
+    "        return self._create_dataloader(self.test_dataset)\n",
+    "\n",
+    "    def get_class_names(self) -> List[str]:\n",
+    "        return self.class_names"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "datamodule = CatDogImageDataModule(\n",
+    "    data_root=\"data\",\n",
+    "    data_dir=\"cats_and_dogs_filtered\",\n",
+    "    batch_size=32,\n",
+    "    num_workers=4,\n",
+    "    train_val_split=[0.8, 0.2],\n",
+    "    pin_memory=True,\n",
+    "    image_size=224,\n",
+    "    url=\"https://download.pytorch.org/tutorials/cats_and_dogs_filtered.zip\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2024-11-10 05:37:17.840\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m81\u001b[0m - \u001b[1mTrain/Validation split: 2241 train, 561 validation images.\u001b[0m\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2024-11-10 05:37:17.910\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36m__main__\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m87\u001b[0m - \u001b[1mTest dataset size: 198 images.\u001b[0m\n"
+     ]
+    }
+   ],
+   "source": [
+    "datamodule.prepare_data()\n",
+    "datamodule.setup()\n",
+    "class_names = datamodule.get_class_names()\n",
+    "train_dataloader = datamodule.train_dataloader()\n",
+    "val_dataloader= datamodule.val_dataloader()\n",
+    "test_dataloader= datamodule.test_dataloader()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['cats', 'dogs']"
+      ]
+     },
+     "execution_count": 36,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "class_names"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

src/datamodules/catdog_datamodule.py

@@ -14,7 +14,8 @@ class CatDogImageDataModule(L.LightningDataModule):
 
     def __init__(
         self,
-        data_dir: Union[str, Path] = "data",
+        data_root: Union[str, Path] = "data",
+        data_dir: Union[str, Path] = "cats_and_dogs_filtered",
         batch_size: int = 32,
         num_workers: int = 4,
         train_val_split: List[float] = [0.8, 0.2],
@@ -23,7 +24,8 @@ class CatDogImageDataModule(L.LightningDataModule):
         url: str = "https://download.pytorch.org/tutorials/cats_and_dogs_filtered.zip",
     ):
         super().__init__()
-        self.data_dir = Path(data_dir)
+        self.data_root = Path(data_root)
+        self.data_dir = data_dir
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.train_val_split = train_val_split
@@ -38,21 +40,27 @@ class CatDogImageDataModule(L.LightningDataModule):
 
     def prepare_data(self):
         """Download the dataset if it doesn't exist."""
-        dataset_path = self.data_dir / "cats_and_dogs_filtered"
-        if not dataset_path.exists():
+        self.dataset_path = self.data_root / self.data_dir
+        if not self.dataset_path.exists():
             logger.info("Downloading and extracting dataset.")
             download_and_extract_archive(
-                url=self.url, download_root=self.data_dir, remove_finished=True
+                url=self.url, download_root=self.data_root, remove_finished=True
             )
             logger.info("Download completed.")
 
     def setup(self, stage: Optional[str] = None):
         """Set up the train, validation, and test datasets."""
 
+        self.prepare_data()
+
         train_transform = transforms.Compose(
             [
                 transforms.Resize((self.image_size, self.image_size)),
-                transforms.RandomHorizontalFlip(),
+                transforms.RandomHorizontalFlip(0.1),
+                transforms.RandomRotation(10),
+                transforms.RandomAffine(0, shear=10, scale=(0.8, 1.2)),
+                transforms.RandomAutocontrast(0.1),
+                transforms.RandomAdjustSharpness(2, 0.1),
                 transforms.ToTensor(),
                 transforms.Normalize(
                     mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
@@ -70,11 +78,12 @@ class CatDogImageDataModule(L.LightningDataModule):
             ]
         )
 
-        train_path = self.data_dir / "cats_and_dogs_filtered" / "train"
-        test_path = self.data_dir / "cats_and_dogs_filtered" / "validation"
+        train_path = self.dataset_path / "train"
+        test_path = self.dataset_path / "test"
 
         if stage == "fit" or stage is None:
             full_train_dataset = ImageFolder(root=train_path, transform=train_transform)
+            self.class_names = full_train_dataset.classes
             train_size = int(self.train_val_split[0] * len(full_train_dataset))
             val_size = len(full_train_dataset) - train_size
             self.train_dataset, self.val_dataset = random_split(
@@ -107,43 +116,42 @@ class CatDogImageDataModule(L.LightningDataModule):
     def test_dataloader(self) -> DataLoader:
         return self._create_dataloader(self.test_dataset)
 
+    def get_class_names(self) -> List[str]:
+        return self.class_names
+
 
 if __name__ == "__main__":
-    from omegaconf import DictConfig, OmegaConf
+    # Test the CatDogImageDataModule
     import hydra
+    from omegaconf import DictConfig, OmegaConf
     import rootutils
 
-    # Setup root directory
-    root = rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
-    logger.info(f"Root directory: {root}")
+    root = rootutils.setup_root(__file__, indicator=".project-root")
 
     @hydra.main(
-        version_base="1.3",
-        config_path=str(root / "configs"),
-        config_name="train",
+        config_path=str(root / "configs"), version_base="1.3", config_name="train"
     )
-    def main(cfg: DictConfig):
-        # Log configuration
-        logger.info("Config:\n" + OmegaConf.to_yaml(cfg))
-
-        # Initialize DataModule
+    def test_datamodule(cfg: DictConfig):
+        logger.info(f"Config:\n{OmegaConf.to_yaml(cfg)}")
         datamodule = CatDogImageDataModule(
+            data_root=cfg.paths.data_dir,
             data_dir=cfg.data.data_dir,
             batch_size=cfg.data.batch_size,
             num_workers=cfg.data.num_workers,
             train_val_split=cfg.data.train_val_split,
             pin_memory=cfg.data.pin_memory,
             image_size=cfg.data.image_size,
-            url=cfg.data.url,
-        )
-        datamodule.prepare_data()
-        datamodule.setup()
-
-        # Log DataLoader sizes
-        logger.info(f"Train DataLoader: {len(datamodule.train_dataloader())} batches")
-        logger.info(
-            f"Validation DataLoader: {len(datamodule.val_dataloader())} batches"
         )
-        logger.info(f"Test DataLoader: {len(datamodule.test_dataloader())} batches")
-
-    main()
+        datamodule.setup(stage="fit")
+        train_loader = datamodule.train_dataloader()
+        val_loader = datamodule.val_dataloader()
+        datamodule.setup(stage="test")
+        test_loader = datamodule.test_dataloader()
+        class_names = datamodule.get_class_names()
+
+        logger.info(f"Train loader: {len(train_loader)} batches")
+        logger.info(f"Validation loader: {len(val_loader)} batches")
+        logger.info(f"Test loader: {len(test_loader)} batches")
+        logger.info(f"Class names: {class_names}")
+
+    test_datamodule()

src/train_new.py

@@ -122,7 +122,7 @@ def run_test_module(
     return test_metrics[0] if test_metrics else {}
 
 
-@hydra.main(config_path="../configs", config_name="train", version_base="1.1")
+@hydra.main(config_path="../configs", config_name="train", version_base="1.3")
 def setup_run_trainer(cfg: DictConfig):
     """Set up and run the Trainer for training and testing."""
     # Display configuration