Merge commit from origin/experiment-runner

Author: kmock930

models/cnn.py

@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, List
 
 import numpy as np
 import torch
@@ -24,6 +24,9 @@
 from .ParamSpace import ParamSpace
 from .base import BaseModel
 
+import numpy as np
+from torch import tensor
+
 MODEL_PATH = Path(".cache/models/cnn_cifar.pth")
 
 
@@ -89,6 +92,7 @@ def __init__(
             nn.Linear(64, num_classes),
         )
 
+    # Will be run by PyTorch under the hood
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.features(x)
         return self.classifier(x)
@@ -104,13 +108,25 @@ def __init__(self, num_classes: int = 10) -> None:
         self._input_channels = 1  # grayscale CIFAR-10
 
     def create_model(self, **params) -> None:
+        """Create the CNN model with given parameters."""        
+        # Extract architecture-specific parameters for Backbone creation      
+        kernel_size = params.get("kernel_size", 3)
+        stride = params.get("stride", 1)
+        learning_rate = params.get("learning_rate", 3e-4)
+        batch_size = params.get("batch_size", 64)
+        weight_decay = params.get("weight_decay", 1e-3)
+        optimizer = params.get("optimizer", "AdamW")
+        
+        # Also ensure default values are set if not provided
+        self.params.setdefault("kernel_size", kernel_size)
+        self.params.setdefault("stride", stride)
+        self.params.setdefault("learning_rate", learning_rate)
+        self.params.setdefault("batch_size", batch_size)
+        self.params.setdefault("weight_decay", weight_decay)
+        self.params.setdefault("optimizer", optimizer)
+
         # Store all parameters passed in
         self.params.update(params)
-        
-        # Extract architecture-specific parameters for Backbone creation
-        kernel_size = self.params.get("kernel_size", 3)
-        stride = self.params.get("stride", 1)
-        
         self.network = Backbone(
             in_channels=self._input_channels,
             num_classes=self.num_classes,
@@ -120,10 +136,9 @@ def create_model(self, **params) -> None:
 
     def train(
         self,
-        X_train: List[np.ndarray],
-        y_train: np.ndarray,
-        X_val: List[np.ndarray],
-        y_val: np.ndarray,
+        train_loader: DataLoader,
+        val_loader: DataLoader,
+        config: Optional[TrainingConfig] = None,
         device: Optional[torch.device] = None,
         epochs: Optional[int] = None,
         patience: Optional[int] = None,
@@ -139,29 +154,7 @@ def train(
         device = device or get_device()
         self.network = self.network.to(device)
 
-        default_config = TrainingConfig()
-        config = TrainingConfig(
-            learning_rate=float(self.params.get("learning_rate", default_config.learning_rate)),
-            weight_decay=float(self.params.get("weight_decay", default_config.weight_decay)),
-            optimizer=self.params.get("optimizer", default_config.optimizer),
-            batch_size=int(self.params.get("batch_size", default_config.batch_size)),
-            # Infrastructure params: use provided values or defaults
-            epochs=epochs if epochs is not None else default_config.epochs,
-            patience=patience if patience is not None else default_config.patience,
-            min_delta=min_delta if min_delta is not None else default_config.min_delta,
-            checkpoint_path=checkpoint_path if checkpoint_path is not None else default_config.checkpoint_path,
-            grad_clip_norm=grad_clip_norm if grad_clip_norm is not None else default_config.grad_clip_norm,
-            writer=writer if writer is not None else default_config.writer,
-        )
-
-        train_loader, val_loader = create_dataloaders(
-            X_train,
-            y_train,
-            X_val,
-            y_val,
-            batch_size=config.batch_size,
-            num_workers=num_workers,
-        )
+        config = config or TrainingConfig()
 
         optimizer = self._build_optimizer(self.network, config)
         scheduler = optim.lr_scheduler.OneCycleLR(
@@ -256,7 +249,8 @@ def train(
 
     def predict(
         self,
-        data_loader: DataLoader,
+        data: DataLoader | List | np.ndarray,
+        labels: List | np.ndarray = None,
         device: Optional[torch.device] = None,
         return_probabilities: bool = False,
     ) -> torch.Tensor:
@@ -266,6 +260,19 @@ def predict(
         network = self.network.to(device)
         network.eval()
 
+        # Handle data input - create DataLoader if raw data is provided
+        if isinstance(data, DataLoader):
+            data_loader = data
+        else:
+            # Raw data provided - create DataLoader
+            from framework.data_utils import create_dataloaders
+            if labels is None:
+                labels = [0] * len(data)  # dummy labels for prediction
+            batch_size = getattr(self, 'params', {}).get('batch_size', 64)
+            _, data_loader = create_dataloaders(
+                data[:1], [labels[0]], data, labels, batch_size=batch_size
+            )
+
         outputs = []
         with torch.no_grad():
             for images, _ in data_loader:
@@ -279,8 +286,7 @@ def predict(
 
     def evaluate(
         self,
-        X: List[np.ndarray],
-        y: np.ndarray,
+        data_loader: DataLoader,
         device: Optional[torch.device] = None,
         criterion: Optional[nn.Module] = None,
         num_workers: int = 0,
@@ -302,6 +308,19 @@ def evaluate(
             pin_memory=torch.cuda.is_available(),
         )
 
+        # Handle data input - create DataLoader if raw data is provided
+        if isinstance(data, DataLoader):
+            data_loader = data
+        else:
+            # Raw data provided - create DataLoader
+            from framework.data_utils import create_dataloaders
+            if labels is None:
+                raise ValueError("labels must be provided when data is not a DataLoader")
+            batch_size = getattr(self, 'params', {}).get('batch_size', 64)
+            _, data_loader = create_dataloaders(
+                data[:1], [labels[0]], data, labels, batch_size=batch_size
+            )
+
         criterion = criterion or nn.CrossEntropyLoss()
 
         total_loss = 0.0

models/decision_tree.py

@@ -70,9 +70,22 @@ def evaluate(self, X_test, y_test) -> Dict[str, float]:
             "precision_weighted": report["weighted avg"]["precision"],
             "recall_weighted": report["weighted avg"]["recall"],
             "f1_weighted": report["weighted avg"]["f1-score"],
-            "roc_auc_weighted": roc_auc_score(y_test, proba, average="weighted", multi_class="ovr"),
         }
 
+        if hasattr(self.estimator, "predict_proba"):
+            try:
+                proba = self.estimator.predict_proba(X_test)
+                if proba.ndim == 2 and proba.shape[1] > 1:
+                    # Check if we have enough classes for ROC AUC calculation
+                    unique_classes = len(set(y_test))
+                    if unique_classes >= 2:
+                        metrics["roc_auc_weighted"] = roc_auc_score(
+                            y_test, proba, average="weighted", multi_class="ovr"
+                        )
+            except (ValueError, Exception) as e:
+                # ROC AUC calculation failed, skip it
+                pass
+
         return metrics
 
     def get_param_space(self) -> Dict[str, ParamSpace]:

models/knn.py

@@ -69,7 +69,22 @@ def evaluate(self, X_test, y_test) -> Dict[str, float]:
             "recall_weighted": report["weighted avg"]["recall"],
             "f1_weighted": report["weighted avg"]["f1-score"],
         }
-        
+
+        # Add ROC AUC if possible
+        if hasattr(self.estimator, "predict_proba"):
+            try:
+                proba = self.estimator.predict_proba(X_test)
+                if proba.ndim == 2 and proba.shape[1] > 1:
+                    # Check if we have enough classes for ROC AUC calculation
+                    unique_classes = len(set(y_test))
+                    if unique_classes >= 2:
+                        metrics["roc_auc_weighted"] = roc_auc_score(
+                            y_test, proba, average="weighted", multi_class="ovr"
+                        )
+            except (ValueError, Exception) as e:
+                # ROC AUC calculation failed, skip it
+                pass
+
         return metrics
 
     def get_param_space(self) -> Dict[str, ParamSpace]: