Merge pull request #4 from AI-Enabled-Software-Testing/pso-and-experiments

Update KNN params, add PSO, create experiment runner

Author: kmock930

framework/data_utils.py

@@ -1,12 +1,11 @@
 """Data loading and preprocessing utilities."""
 
 from pathlib import Path
-from typing import List, Tuple
+from typing import Any, Dict, List, Tuple
 import numpy as np
 from datasets import load_from_disk
 from sklearn.model_selection import train_test_split
 from torch.utils.data import DataLoader
-from PIL import Image
 
 from framework import utils
 from framework.datasets import CIFAR10Dataset
@@ -116,14 +115,12 @@ def create_dataloaders(
     X_val: List[np.ndarray],
     y_val: np.ndarray,
     batch_size: int,
-    num_workers: int = 2,
 ) -> Tuple[DataLoader, DataLoader]:
     train_dataset = CIFAR10Dataset(X_train, y_train)
     train_loader = DataLoader(
         train_dataset,
         batch_size=batch_size,
         shuffle=True,
-        num_workers=num_workers,
         pin_memory=utils.is_cuda_available(),
     )
 
@@ -132,8 +129,38 @@ def create_dataloaders(
         val_dataset,
         batch_size=batch_size,
         shuffle=False,
-        num_workers=num_workers,
         pin_memory=utils.is_cuda_available(),
     )
 
     return train_loader, val_loader
+
+
+def prepare_dataset(val_ratio: float = 0.1) -> Dict[str, Any]:
+    """Prepare and return the CIFAR-10 dataset"""
+    ds_dict = load_cifar10_data()
+    train_images, train_labels = prepare_data(ds_dict, "train")
+    test_images, test_labels = prepare_data(ds_dict, "test")
+
+    X_train, y_train, X_val, y_val = split_train_val(
+        train_images, train_labels, val_ratio=val_ratio
+    )
+
+    def flatten(images):
+        stacked = np.stack([np.asarray(img, dtype=np.float32) for img in images])
+        return stacked.reshape(len(images), -1)
+
+    train_flat = flatten(X_train)
+    val_flat = flatten(X_val)
+    test_flat = flatten(test_images)
+
+    return {
+        "train_images": X_train,
+        "train_labels": y_train,
+        "val_images": X_val,
+        "val_labels": y_val,
+        "test_images": test_images,
+        "test_labels": test_labels,
+        "train_flat": train_flat,
+        "val_flat": val_flat,
+        "test_flat": test_flat,
+    }

framework/fitness.py

@@ -1,12 +1,12 @@
 def calculate_composite_fitness(metrics: dict[str, float]) -> float:
     """Calculate composite fitness score from evaluation metrics."""
     # Extract metrics
-    f1_macro = metrics.get("f1_macro", 0.0)
-    recall_macro = metrics.get("recall_macro", 0.0)
-    roc_auc = metrics.get("roc_auc", 0.0)
-    precision_macro = metrics.get("precision_macro", 0.0)
-    accuracy = metrics.get("accuracy", 0.0)
-    f1_micro = metrics.get("f1_micro", 0.0)
+    f1_macro = metrics["f1_macro"]
+    recall_macro = metrics["recall_macro"]
+    roc_auc = metrics["roc_auc"]
+    precision_macro = metrics["precision_macro"]
+    accuracy = metrics["accuracy"]
+    f1_micro = metrics["f1_micro"]
 
     # Composite fitness
     composite_fitness = (

hparam_search.py

@@ -8,13 +8,12 @@
 import torch
 from torch.utils.tensorboard import SummaryWriter
 
-from framework.data_utils import (
-    load_cifar10_data,
-    prepare_data,
-    split_train_val,
-)
+from framework.data_utils import prepare_dataset
 from framework.fitness import calculate_composite_fitness
-from models.base import get_model_by_name
+from models.cnn import CNNModel
+from models.decision_tree import DecisionTreeModel
+from models.factory import get_model_by_name
+from models.knn import KNNModel
 from search import RandomSearch
 
 RANDOM_SEED = 321
@@ -34,34 +33,6 @@ def set_seeds(seed: int):
         torch.cuda.manual_seed_all(seed)
 
 
-def prepare_dataset() -> Dict[str, Any]:
-    ds_dict = load_cifar10_data()
-    train_images, train_labels = prepare_data(ds_dict, "train")
-    test_images, test_labels = prepare_data(ds_dict, "test")
-
-    X_train, y_train, X_val, y_val = split_train_val(
-        train_images, train_labels, val_ratio=0.2
-    )
-
-    def flatten(images):
-        stacked = np.stack([np.asarray(img, dtype=np.float32) for img in images])
-        return stacked.reshape(len(images), -1)
-
-    train_flat = flatten(X_train)
-    val_flat = flatten(X_val)
-    test_flat = flatten(test_images)
-
-    return {
-        "train_images": X_train,
-        "train_labels": y_train,
-        "val_images": X_val,
-        "val_labels": y_val,
-        "test_images": test_images,
-        "test_labels": test_labels,
-        "train_flat": train_flat,
-        "val_flat": val_flat,
-        "test_flat": test_flat,
-    }
 
 
 def evaluate_model(
@@ -72,10 +43,12 @@ def evaluate_model(
     model = get_model_by_name(model_key)
 
     if model_key in {"dt", "knn"}:
+        assert isinstance(model, (DecisionTreeModel, KNNModel))
         model.create_model(**params)
         model.train(data["train_flat"], data["train_labels"])
         metrics = model.evaluate(data["val_flat"], data["val_labels"])
     elif model_key == "cnn":
+        assert isinstance(model, CNNModel)
         model.create_model(**params)
         model.train(
             data["train_images"],

models/base.py

@@ -1,7 +1,9 @@
 """Abstract interface for models used in the hyperparameter tuning framework."""
 
 from abc import ABC, abstractmethod
-from typing import Dict, Any, Literal
+from typing import Dict, Any
+
+
 
 from .ParamSpace import ParamSpace
 
@@ -36,23 +38,3 @@ def evaluate(self, *args: Any, **kwargs: Any) -> Dict[str, float]:
     def get_param_space(self) -> Dict[str, ParamSpace]:
         """Return the searchable hyperparameter space."""
         raise NotImplementedError
-
-
-def get_model_by_name(model_name: Literal["dt", "knn", "cnn"]) -> BaseModel:
-    """Factory function to get model by name."""
-    from models.decision_tree import DecisionTreeModel
-    from models.knn import KNNModel
-    from models.cnn import CNNModel
-
-    models = {
-        "dt": DecisionTreeModel,
-        "knn": KNNModel,
-        "cnn": CNNModel,
-    }
-
-    if model_name not in models:
-        raise ValueError(
-            f"Unknown model: {model_name}. Available models: {list(models.keys())}"
-        )
-
-    return models[model_name]()