Fix: notebook synced with src, added CNNDropoutModel, cleanup config

Author: afarresg

README.md

@@ -30,13 +30,15 @@ Clasificador de pistachos (Kirmizi vs Siirt) usando Deep Learning con PyTorch Li
 
 ## Configuracion
 
-Editar `config/configuracion.yaml`:
+Editar `config/configuracion.yaml` (solo hiperparametros del modelo):
 
 - `model_type`: `cnn_batchnorm` o `cnn_dropout`
 - `learning_rate`, `batch_size`, `max_epochs`, `patience`
-- `image_size` (default: 28)
+- `dropout_rate`
 - `wandb_project`: nombre del proyecto en W&B
 
+El resto (`output_size=2`, `image_size=28`, splits, etc.) son constantes de arquitectura en el codigo.
+
 ## Ejecucion local
 
 ```bash
@@ -48,6 +50,9 @@ python src/main.py configuracion.yaml
 # Sin W&B (offline)
 python src/main.py configuracion.yaml --no-wandb
 
+# Smoke test (10% datos, 3 epocas)
+python src/main.py configuracion.yaml --no-wandb --data-fraction 0.1
+
 # W&B Sweep (grid search)
 python src/main.py configuracion.yaml --sweep
 

config/configuracion.yaml

@@ -1,22 +1,11 @@
-# Parametros de entrenamiento
+# Hiperparametros del modelo (se modifican via sweep)
 semilla: 42
 max_epochs: 50
 patience: 5
 batch_size: 64
 learning_rate: 0.001
-
-# Parametros del modelo
 model_type: "cnn_batchnorm"
-output_size: 2
 dropout_rate: 0.3
 
-# Parametros de datos
-image_size: 28
-data_fraction: 1.0
-val_split: 0.15
-test_split: 0.15
-num_workers: 4
-data_dir: "./data"
-
 # W&B
 wandb_project: "pistachio-mlops"

notebook/pistachio.ipynb

@@ -93,24 +93,27 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import logging\n",
-    "import os\n",
-    "import shutil\n",
-    "from pathlib import Path\n",
-    "\n",
-    "import kagglehub\n",
-    "import lightning as pl\n",
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
-    "import torch\n",
-    "import torch.nn as nn\n",
-    "import torch.optim as optim\n",
-    "import torchvision.transforms as transforms\n",
-    "from lightning import Trainer\n",
-    "from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint\n",
-    "from sklearn.metrics import classification_report, confusion_matrix, f1_score\n",
-    "from torch.utils.data import DataLoader, random_split\n",
-    "from torchvision.datasets import ImageFolder"
+"import logging\n",
+"import os\n",
+"import shutil\n",
+"from pathlib import Path\n",
+"\n",
+"import kagglehub\n",
+"import lightning as pl\n",
+"import matplotlib.pyplot as plt\n",
+"import numpy as np\n",
+"import torch\n",
+"import torch.nn as nn\n",
+"import torch.optim as optim\n",
+"import torchvision.transforms as transforms\n",
+"from lightning import Trainer\n",
+"from lightning.pytorch.callbacks import EarlyStopping, ModelCheckpoint\n",
+"from sklearn.metrics import classification_report, confusion_matrix, f1_score\n",
+"from torch.utils.data import DataLoader, random_split\n",
+"from torchvision.datasets import ImageFolder\n",
+"\n",
+"logging.basicConfig(level=logging.INFO)\n",
+"logger = logging.getLogger(__name__)"
    ]
   },
   {
@@ -127,17 +130,18 @@
    "id": "a0005",
    "metadata": {},
    "outputs": [],
-   "source": [
+    "source": [
     "class PistachioDataModule(pl.LightningDataModule):\n",
     "    def __init__(\n",
     "        self,\n",
     "        batch_size=64,\n",
-    "        num_workers=8,\n",
-    "        data_dir=\"./pistachio_data\",\n",
+    "        num_workers=4,\n",
+    "        data_dir=\"./data\",\n",
     "        data_fraction=1.0,\n",
     "        val_split=0.15,\n",
     "        test_split=0.15,\n",
     "        image_size=28,\n",
+    "        seed=42,\n",
     "    ):\n",
     "        super().__init__()\n",
     "        self.batch_size = batch_size\n",
@@ -147,18 +151,20 @@
     "        self.val_split = val_split\n",
     "        self.test_split = test_split\n",
     "        self.image_size = image_size\n",
+    "        self.seed = seed\n",
     "        self.dataset_path = None\n",
+    "        self.class_names = None\n",
     "\n",
     "    def prepare_data(self):\n",
     "        if not os.path.exists(self.data_dir) or not any(Path(self.data_dir).iterdir()):\n",
-    "            print(\"Downloading pistachio dataset from Kaggle...\")\n",
+    "            logger.info(\"Descargando dataset de pistachos desde Kaggle...\")\n",
     "            raw_path = kagglehub.dataset_download(\n",
     "                \"muratkokludataset/pistachio-image-dataset\"\n",
     "            )\n",
     "            self.dataset_path = os.path.join(\n",
     "                raw_path, \"Pistachio_Image_Dataset\", \"Pistachio_Image_Dataset\"\n",
     "            )\n",
-    "            print(f\"Dataset downloaded to: {raw_path}\")\n",
+    "            logger.info(\"Dataset descargado en: %s\", raw_path)\n",
     "\n",
     "            os.makedirs(self.data_dir, exist_ok=True)\n",
     "            src = Path(self.dataset_path)\n",
@@ -168,7 +174,7 @@
     "                if item.is_dir() and not dest.exists():\n",
     "                    shutil.copytree(item, dest)\n",
     "\n",
-    "            print(f\"Dataset prepared in: {self.data_dir}\")\n",
+    "            logger.info(\"Dataset copiado a: %s\", self.data_dir)\n",
     "\n",
     "    def setup(self, stage=None):\n",
     "        transform = transforms.Compose(\n",
@@ -180,22 +186,31 @@
     "        )\n",
     "\n",
     "        full_dataset = ImageFolder(self.data_dir, transform=transform)\n",
+    "        self.class_names = full_dataset.classes\n",
     "        dataset_size = len(full_dataset)\n",
     "\n",
+    "        if self.data_fraction < 1.0:\n",
+    "            subset_size = int(dataset_size * self.data_fraction)\n",
+    "            full_dataset, _ = random_split(\n",
+    "                full_dataset,\n",
+    "                [subset_size, dataset_size - subset_size],\n",
+    "                generator=torch.Generator().manual_seed(self.seed),\n",
+    "            )\n",
+    "            dataset_size = subset_size\n",
+    "\n",
     "        test_size = int(dataset_size * self.test_split)\n",
     "        val_size = int(dataset_size * self.val_split)\n",
     "        train_size = dataset_size - val_size - test_size\n",
     "\n",
     "        self.train_dataset, self.val_dataset, self.test_dataset = random_split(\n",
     "            full_dataset,\n",
     "            [train_size, val_size, test_size],\n",
-    "            generator=torch.Generator().manual_seed(42),\n",
+    "            generator=torch.Generator().manual_seed(self.seed),\n",
     "        )\n",
     "\n",
-    "        print(f\"Classes: {full_dataset.classes}\")\n",
-    "        print(f\"Train size: {len(self.train_dataset)}\")\n",
-    "        print(f\"Val size: {len(self.val_dataset)}\")\n",
-    "        print(f\"Test size: {len(self.test_dataset)}\")\n",
+    "        logger.info(\"Clases: %s\", self.class_names)\n",
+    "        logger.info(\"Train: %d, Val: %d, Test: %d\",\n",
+    "                     train_size, val_size, test_size)\n",
     "\n",
     "    def train_dataloader(self):\n",
     "        return DataLoader(\n",
@@ -265,7 +280,7 @@
     "dm.prepare_data()\n",
     "dm.setup()\n",
     "\n",
-    "class_names = dm.train_dataset.dataset.classes\n",
+    "class_names = dm.class_names\n",
     "\n",
     "fig, axes = plt.subplots(2, 4, figsize=(12, 6))\n",
     "fig.suptitle(\"Samples\", fontsize=14)\n",
@@ -382,7 +397,7 @@
     "        self.log(\"test_acc\", acc, on_step=False, on_epoch=True, prog_bar=True)\n",
     "        self.test_predictions.extend(predicted.cpu().numpy())\n",
     "        self.test_labels.extend(y_batch.cpu().numpy())\n",
-    "        return acc\n",
+    "        return {\"loss\": loss, \"acc\": acc}\n",
     "\n",
     "    def on_test_epoch_start(self):\n",
     "        self.test_predictions = []\n",
@@ -415,10 +430,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "class CNNBatchNormModel(BaselineModel):\n",
-    "    def __init__(self, output_size=2, lr=0.001):\n",
-    "        super().__init__(output_size=output_size, lr=lr)\n",
-    "        self.save_hyperparameters()\n",
+"class CNNBatchNormModel(BaselineModel):\n",
+"    def __init__(self, output_size=2, lr=0.001, dropout_rate=0.3):\n",
+"        super().__init__(output_size=output_size, lr=lr)\n",
+"        self.save_hyperparameters()\n",
     "        self.features = nn.Sequential(\n",
     "            nn.Conv2d(3, 32, kernel_size=3, padding=1, bias=False),\n",
     "            nn.BatchNorm2d(32),\n",
@@ -440,16 +455,53 @@
     "            nn.Linear(64 * 7 * 7, 128),\n",
     "            nn.BatchNorm1d(128),\n",
     "            nn.ReLU(),\n",
-    "            nn.Dropout(0.3),\n",
-    "            nn.Linear(128, output_size),\n",
-    "        )\n",
-    "\n",
-    "    def forward(self, x):\n",
-    "        x = self.features(x)\n",
-    "        return self.classifier(x)\n",
-    "\n",
-    "    def configure_optimizers(self):\n",
-    "        return optim.Adam(self.parameters(), lr=self.hparams.lr)"
+"            nn.Dropout(dropout_rate),\n",
+"            nn.Linear(128, output_size),\n",
+"        )\n",
+"\n",
+"    def forward(self, x):\n",
+"        x = self.features(x)\n",
+"        return self.classifier(x)\n",
+"\n",
+"    def configure_optimizers(self):\n",
+"        return optim.Adam(self.parameters(), lr=self.hparams.lr)\n",
+"\n",
+"\n",
+"class CNNDropoutModel(BaselineModel):\n",
+"    \"\"\"CNN con Dropout en lugar de BatchNorm para comparacion.\"\"\"\n",
+"    def __init__(self, output_size=2, lr=0.001, dropout_rate=0.3):\n",
+"        super().__init__(output_size=output_size, lr=lr)\n",
+"        self.save_hyperparameters()\n",
+"        self.features = nn.Sequential(\n",
+"            nn.Conv2d(3, 32, kernel_size=3, padding=1, bias=False),\n",
+"            nn.ReLU(),\n",
+"            nn.Dropout2d(dropout_rate),\n",
+"            nn.Conv2d(32, 32, kernel_size=3, padding=1, bias=False),\n",
+"            nn.ReLU(),\n",
+"            nn.Dropout2d(dropout_rate),\n",
+"            nn.MaxPool2d(2),\n",
+"            nn.Conv2d(32, 64, kernel_size=3, padding=1, bias=False),\n",
+"            nn.ReLU(),\n",
+"            nn.Dropout2d(dropout_rate),\n",
+"            nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=False),\n",
+"            nn.ReLU(),\n",
+"            nn.Dropout2d(dropout_rate),\n",
+"            nn.MaxPool2d(2),\n",
+"        )\n",
+"        self.classifier = nn.Sequential(\n",
+"            nn.Flatten(),\n",
+"            nn.Linear(64 * 7 * 7, 128),\n",
+"            nn.ReLU(),\n",
+"            nn.Dropout(dropout_rate),\n",
+"            nn.Linear(128, output_size),\n",
+"        )\n",
+"\n",
+"    def forward(self, x):\n",
+"        x = self.features(x)\n",
+"        return self.classifier(x)\n",
+"\n",
+"    def configure_optimizers(self):\n",
+"        return optim.Adam(self.parameters(), lr=self.hparams.lr)"
    ]
   },
   {
@@ -707,11 +759,17 @@
      ]
     }
    ],
-   "source": [
+    "source": [
     "data_module = PistachioDataModule(data_fraction=1.0)\n",
     "lrs = [0.01]\n",
-    "results = train_models(CNNBatchNormModel, lrs, data_module)\n",
-    "print_model_results(results)"
+    "\n",
+    "logging.info(\"Entrenando CNNBatchNormModel...\")\n",
+    "results_bn = train_models(CNNBatchNormModel, lrs, data_module)\n",
+    "print_model_results(results_bn)\n",
+    "\n",
+    "logging.info(\"Entrenando CNNDropoutModel...\")\n",
+    "results_drop = train_models(CNNDropoutModel, lrs, data_module)\n",
+    "print_model_results(results_drop)"
    ]
   }
  ],

src/api_inferencia.py

@@ -9,6 +9,10 @@
 from src.model import get_model
 from src.utils import get_project_root, load_config
 
+OUTPUT_SIZE = 2
+IMAGE_SIZE = 28
+CLASS_NAMES = ["Kirmizi_Pistachio", "Siirt_Pistachio"]
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -18,8 +22,8 @@ async def lifespan(app: FastAPI):
     model_type = parametros.get("model_type", "cnn_batchnorm")
     model = get_model(
         model_type=model_type,
-        output_size=int(parametros["output_size"]),
-        lr=float(parametros.get("learning_rate", 0.001)),
+        output_size=OUTPUT_SIZE,
+        lr=0.001,
         dropout_rate=float(parametros.get("dropout_rate", 0.3)),
     )
 
@@ -29,8 +33,8 @@ async def lifespan(app: FastAPI):
     model.eval()
 
     app.state.model = model
-    app.state.class_names = ["Kirmizi_Pistachio", "Siirt_Pistachio"]
-    app.state.image_size = int(parametros.get("image_size", 28))
+    app.state.class_names = CLASS_NAMES
+    app.state.image_size = IMAGE_SIZE
 
     yield
 

src/main.py

@@ -37,6 +37,12 @@ def main():
         action="store_true",
         help="Desactivar logging a W&B",
     )
+    parser.add_argument(
+        "--data-fraction",
+        type=float,
+        default=1.0,
+        help="Fraccion de datos a usar (0-1, util para smoke tests)",
+    )
     args = parser.parse_args()
 
     setup_logging("info")
@@ -52,7 +58,7 @@ def main():
             count=args.count,
         )
     else:
-        train_model(config, wandb_log=not args.no_wandb)
+        train_model(config, wandb_log=not args.no_wandb, data_fraction=args.data_fraction)
 
 
 if __name__ == "__main__":

src/sweep.py

@@ -16,7 +16,7 @@
         "learning_rate": {"values": [0.001, 0.01]},
         "batch_size": {"values": [32, 64]},
         "model_type": {"values": ["cnn_batchnorm", "cnn_dropout"]},
-        "image_size": {"values": [28]},
+        "dropout_rate": {"values": [0.3]},
     },
 }