model trainer fixed

config/config.yaml

@@ -14,5 +14,10 @@ data_transformation:
 
 model_building:
   root_dir : artifacts/model
+
+model_trainer:
+  root_dir : artifacts/trained_model
+  test_data_path : artifacts/data_transformation/test_dataset.pt
+  train_data_path : artifacts/data_transformation/train_dataset.pt
   
   

lightning_logs/version_0/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_1/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_10/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_11/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_12/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_13/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_14/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_15/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_16/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_17/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_18/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_19/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_2/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_3/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_4/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_5/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_6/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_7/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_8/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

lightning_logs/version_9/hparams.yaml

@@ -0,0 +1,7 @@
+display_step: 10
+in_channels: 1
+lambda_gp: 10
+lambda_r1: 10
+lambda_recon: 100
+learning_rate: 0.0002
+out_channels: 2

params.yaml

@@ -23,3 +23,11 @@ OUTPUT_CHANNELS: 2
 
 # Critic Parameters
 IN_CHANNELS: 3
+
+# model train
+LEARNING_RATE : 2e-4
+LAMBDA_RECON : 100
+DISPLAY_STEP : 10
+INPUT_CHANNELS : 1
+OUTPUT_CHANNELS : 2
+EPOCH : 1

research/data_transformation.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
@@ -21,7 +21,7 @@
        "'c:\\\\mlops project\\\\image-colorization-mlops'"
       ]
      },
-     "execution_count": 2,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -32,7 +32,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -51,7 +51,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -88,13 +88,13 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
     "import numpy as np\n",
     "import torch\n",
-    "from torch.utils.data import DataLoader, Dataset\n",
+    "from torch.utils.data import Dataset\n",
     "from torchvision import transforms\n",
     "\n",
     "class ImageColorizationDataset(Dataset):\n",
@@ -106,24 +106,26 @@
     "        return len(self.dataset[0])\n",
     "    \n",
     "    def __getitem__(self, idx):\n",
-    "        L = np.array(self.dataset[0][idx]).reshape(self.image_size + (1,))\n",
+    "        L = np.array(self.dataset[0][idx]).reshape(self.image_size)\n",
     "        L = transforms.ToTensor()(L)\n",
     "        \n",
     "        ab = np.array(self.dataset[1][idx])\n",
     "        ab = transforms.ToTensor()(ab)\n",
     "        \n",
-    "        return ab, L\n"
+    "        return ab, L"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
     "from torch.utils.data import DataLoader\n",
     "import gc\n",
     "import os\n",
+    "import numpy as np\n",
+    "import torch\n",
     "from src.imagecolorization.logging import logger\n",
     "\n",
     "class DataTransformation:\n",
@@ -137,71 +139,51 @@
     "        gc.collect()\n",
     "        return dataset\n",
     "    \n",
-    "    \n",
-    "    def get_dataloader(self, dataset):\n",
+    "    def get_datasets(self, dataset):\n",
     "        train_dataset = ImageColorizationDataset(\n",
     "            dataset=dataset,\n",
-    "            image_size=self.config.IMAGE_SIZE\n",
     "        )\n",
     "        test_dataset = ImageColorizationDataset(\n",
     "            dataset=dataset,\n",
-    "            image_size=self.config.IMAGE_SIZE\n",
-    "        )\n",
-    "        \n",
-    "        train_loader = DataLoader(\n",
-    "            train_dataset,\n",
-    "            batch_size=self.config.BATCH_SIZE,\n",
-    "            shuffle=True,\n",
-    "            pin_memory=True\n",
     "        )\n",
-    "        test_loader = DataLoader(\n",
-    "            test_dataset,\n",
-    "            batch_size=self.config.BATCH_SIZE,\n",
-    "            shuffle=True,\n",
-    "            pin_memory=True\n",
-    "        )\n",
-    "        \n",
-    "        return train_loader, test_loader\n",
-    "    \n",
-    "    \n",
     "    \n",
+    "        return train_dataset, test_dataset\n",
     "        \n",
     "        \n",
     "    \n",
-    "    \n",
-    "    def save_dataloaders(self, train_loader, test_loader):\n",
+    "    def save_datasets(self, train_dataset, test_dataset):\n",
     "        # Ensure the directory exists\n",
     "        os.makedirs(self.config.root_dir, exist_ok=True)\n",
     "\n",
-    "        train_loader_path = os.path.join(self.config.root_dir, 'train_loader.pt')\n",
-    "        test_loader_path = os.path.join(self.config.root_dir, 'test_loader.pt')\n",
+    "        train_dataset_path = os.path.join(self.config.root_dir, 'train_dataset.pt')\n",
+    "        test_dataset_path = os.path.join(self.config.root_dir, 'test_dataset.pt')\n",
     "\n",
     "        try:\n",
-    "            # Save the dataloaders\n",
-    "            torch.save(train_loader, train_loader_path)\n",
-    "            torch.save(test_loader, test_loader_path)\n",
+    "            # Save the datasets\n",
+    "            torch.save(train_dataset, train_dataset_path)\n",
+    "            torch.save(test_dataset, test_dataset_path)\n",
     "\n",
-    "            logger.info(f\"Train Loader saved at: {train_loader_path}\")\n",
-    "            logger.info(f\"Test Loader saved at: {test_loader_path}\")\n",
+    "            logger.info(f\"Train dataset saved at: {train_dataset_path}\")\n",
+    "            logger.info(f\"Test dataset saved at: {test_dataset_path}\")\n",
     "        except Exception as e:\n",
-    "            logger.error(f\"Error saving dataloaders: {str(e)}\")\n",
-    "            raise e\n"
+    "            logger.error(f\"Error saving datasets: {str(e)}\")\n",
+    "            raise e"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 12,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "[2024-08-18 17:50:45,127: INFO: common: yaml file: config\\config.yaml loaded successfully]\n",
-      "[2024-08-18 17:50:45,129: INFO: common: yaml file: params.yaml loaded successfully]\n",
-      "[2024-08-18 17:50:45,129: INFO: common: created directory at: artifacts]\n",
-      "[2024-08-18 17:50:57,600: INFO: 2567581832: Train Loader saved at: artifacts/data_transformation\\train_loader.pt]\n",
-      "[2024-08-18 17:50:57,605: INFO: 2567581832: Test Loader saved at: artifacts/data_transformation\\test_loader.pt]\n"
+      "[2024-08-24 19:09:25,021: INFO: common: yaml file: config\\config.yaml loaded successfully]\n",
+      "[2024-08-24 19:09:25,024: INFO: common: yaml file: params.yaml loaded successfully]\n",
+      "[2024-08-24 19:09:25,026: INFO: common: created directory at: artifacts]\n",
+      "[2024-08-24 19:09:43,417: INFO: 3400243030: Train dataset saved at: artifacts/data_transformation\\train_dataset.pt]\n",
+      "[2024-08-24 19:09:43,440: INFO: 3400243030: Test dataset saved at: artifacts/data_transformation\\test_dataset.pt]\n"
      ]
     }
    ],
@@ -215,10 +197,10 @@
     "    dataset = data_transformation.load_data()\n",
     "    \n",
     "    # Get the dataloader using the loaded dataset\n",
-    "    train_loader, test_loader = data_transformation.get_dataloader(dataset)\n",
+    "    train_dataset, test_dataset = data_transformation.get_datasets(dataset)\n",
     "    \n",
-    "    # Perform any further operations (e.g., saving the dataloaders)\n",
-    "    data_transformation.save_dataloaders(train_loader, test_loader)\n",
+    "    # Perform any further operations (e.g., saving the dataset)\n",
+    "    data_transformation.save_datasets(train_dataset, test_dataset)\n",
     "    \n",
     "except Exception as e:\n",
     "    raise e\n"

src/imagecolorization/conponents/data_tranformation.py

@@ -18,7 +18,7 @@ class ImageColorizationDataset(Dataset):
         return len(self.dataset[0])
     
     def __getitem__(self, idx):
-        L = np.array(self.dataset[0][idx]).reshape(self.image_size + (1,))
+        L = np.array(self.dataset[0][idx]).reshape(self.image_size)
         L = transforms.ToTensor()(L)
         
         ab = np.array(self.dataset[1][idx])
@@ -27,6 +27,13 @@ class ImageColorizationDataset(Dataset):
         return ab, L
     
     
+from torch.utils.data import DataLoader
+import gc
+import os
+import numpy as np
+import torch
+from src.imagecolorization.logging import logger
+
 class DataTransformation:
     def __init__(self, config: DataTransformationConfig):
         self.config = config
@@ -38,8 +45,7 @@ class DataTransformation:
         gc.collect()
         return dataset
     
-    
-    def get_dataloader(self, dataset):
+    def get_datasets(self, dataset):
         train_dataset = ImageColorizationDataset(
             dataset=dataset,
             image_size=self.config.IMAGE_SIZE
@@ -49,21 +55,25 @@ class DataTransformation:
             image_size=self.config.IMAGE_SIZE
         )
         
-        train_loader = DataLoader(
-            train_dataset,
-            batch_size=self.config.BATCH_SIZE,
-            shuffle=True,
-            pin_memory=True
-        )
-        test_loader = DataLoader(
-            test_dataset,
-            batch_size=self.config.BATCH_SIZE,
-            shuffle=True,
-            pin_memory=True
-        )
-        
-        return train_loader, test_loader
+        return train_dataset, test_dataset
     
+    def save_datasets(self, train_dataset, test_dataset):
+        # Ensure the directory exists
+        os.makedirs(self.config.root_dir, exist_ok=True)
+
+        train_dataset_path = os.path.join(self.config.root_dir, 'train_dataset.pt')
+        test_dataset_path = os.path.join(self.config.root_dir, 'test_dataset.pt')
+
+        try:
+            # Save the datasets
+            torch.save(train_dataset, train_dataset_path)
+            torch.save(test_dataset, test_dataset_path)
+
+            logger.info(f"Train dataset saved at: {train_dataset_path}")
+            logger.info(f"Test dataset saved at: {test_dataset_path}")
+        except Exception as e:
+            logger.error(f"Error saving datasets: {str(e)}")
+            raise e
     
     
         
@@ -87,4 +97,3 @@ class DataTransformation:
         except Exception as e:
             logger.error(f"Error saving dataloaders: {str(e)}")
             raise e
-    

src/imagecolorization/pipeline/stage02_data_transformation.py

@@ -10,6 +10,12 @@ class DataTransformationPipeline:
         config = ConfigurationManager()
         data_transformation_config = config.get_data_transformation_config()
         data_transformation = DataTransformation(config=data_transformation_config)
+    
+        # Load the dataset
         dataset = data_transformation.load_data()
-        train_loader, test_loader = data_transformation.get_dataloader(dataset)
-        data_transformation.save_dataloaders(train_loader, test_loader)
+        
+        # Get the dataloader using the loaded dataset
+        train_dataset, test_dataset = data_transformation.get_datasets(dataset)
+    
+        # Perform any further operations (e.g., saving the dataset)
+        data_transformation.save_datasets(train_dataset, test_dataset)