Merge pull request #3 from AI-Enabled-Software-Testing/experiment-runner

Add experiment runner and composite fitness function

Author: seofernando25

framework/fitness.py

@@ -0,0 +1,22 @@
+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)
+    
+    # Composite fitness
+    composite_fitness = (
+        0.30 * f1_macro +
+        0.20 * recall_macro +
+        0.20 * roc_auc +
+        0.15 * precision_macro +
+        0.10 * accuracy +
+        0.05 * f1_micro
+    )
+    
+    return composite_fitness
+

framework/training.py

@@ -76,16 +76,16 @@ def train_epoch(
     epoch: int = 0,
     grad_clip_norm: float = 1.0,
     writer: Optional[SummaryWriter] = None,
+    verbose: bool = True,
 ) -> Tuple[float, float]:
     """Trains the model for one epoch and returns the epoch loss and accuracy."""
     nn.Module.train(model, mode=True)
     running_loss = 0.0
     correct = 0
     total = 0
 
-    for batch_idx, (images, labels) in enumerate(
-        tqdm(train_loader, desc="Training", leave=False)
-    ):
+    loader = tqdm(train_loader, desc="Training", leave=False) if verbose else train_loader
+    for batch_idx, (images, labels) in enumerate(loader):
         images, labels = images.to(device), labels.to(device)
 
         optimizer.zero_grad()
@@ -131,6 +131,7 @@ def validate(
     device: torch.device,
     epoch: int = 0,
     writer: Optional[SummaryWriter] = None,
+    verbose: bool = True,
 ) -> Tuple[float, float]:
     """Validates the model and returns the epoch loss and accuracy."""
     model.eval()
@@ -139,7 +140,8 @@ def validate(
     total = 0
 
     with torch.no_grad():
-        for images, labels in tqdm(val_loader, desc="Validating", leave=False):
+        loader = tqdm(val_loader, desc="Validating", leave=False) if verbose else val_loader
+        for images, labels in loader:
             images, labels = images.to(device), labels.to(device)
 
             outputs = model(images)

hparam_search.py

@@ -13,6 +13,7 @@
     prepare_data,
     split_train_val,
 )
+from framework.fitness import calculate_composite_fitness
 from models.base import get_model_by_name
 from search import RandomSearch
 
@@ -73,19 +74,22 @@ def evaluate_model(
     if model_key in {"dt", "knn"}:
         model.create_model(**params)
         model.train(data["train_flat"], data["train_labels"])
-        return model.evaluate(data["val_flat"], data["val_labels"])
-
-    if model_key == "cnn":
+        metrics = model.evaluate(data["val_flat"], data["val_labels"])
+    elif model_key == "cnn":
         model.create_model(**params)
         model.train(
             data["train_images"],
             data["train_labels"],
             data["val_images"],
             data["val_labels"],
         )
-        return model.evaluate(data["val_images"], data["val_labels"])
-
-    raise ValueError(f"Unsupported model key: {model_key}")
+        metrics = model.evaluate(data["val_images"], data["val_labels"])
+    else:
+        raise ValueError(f"Unsupported model key: {model_key}")
+    
+    metrics["composite_fitness"] = calculate_composite_fitness(metrics)
+    
+    return metrics
 
 
 def run_search(model_key: Literal["dt", "knn", "cnn"], trials: int) -> None:
@@ -98,7 +102,7 @@ def run_search(model_key: Literal["dt", "knn", "cnn"], trials: int) -> None:
     searcher = RandomSearch(
         param_space=param_space,
         evaluate_fn=lambda sampled: evaluate_model(model_key, sampled, data),
-        metric_key="accuracy",
+        metric_key="composite_fitness",
         seed=RANDOM_SEED,
     )
 
@@ -113,7 +117,11 @@ def run_search(model_key: Literal["dt", "knn", "cnn"], trials: int) -> None:
     print("-" * 80)
     print(f"Model: {model_key}")
     print(f"Trials: {trials}")
-    print(f"Best val accuracy: {result.best_metrics['accuracy']:.4f}")
+    print(f"Best composite fitness: {result.best_metrics['composite_fitness']:.4f}")
+    print("Best metrics:")
+    for name, value in result.best_metrics.items():
+        if isinstance(value, float):
+            print(f"  {name}: {value:.4f}")
     print("Best hyperparameters:")
     for name, value in result.best_params.items():
         print(f"  {name}: {value}")

models/ParamSpace.py

@@ -6,6 +6,7 @@
 class ParamType(Enum):
     INTEGER = "integer"
     FLOAT = "float"
+    FLOAT_LOG = "float_log"
     CATEGORICAL = "categorical"
     BOOLEAN = "boolean"
 
@@ -20,11 +21,16 @@ class ParamSpace:
 
     def __post_init__(self):
         """Validate parameter space configuration"""
-        if self.param_type in [ParamType.INTEGER, ParamType.FLOAT]:
+        if self.param_type in [ParamType.INTEGER, ParamType.FLOAT, ParamType.FLOAT_LOG]:
             if self.min_value is None or self.max_value is None:
                 raise ValueError(
                     f"min_value and max_value required for {self.param_type.value}"
                 )
+            if self.param_type == ParamType.FLOAT_LOG:
+                if float(self.min_value) <= 0 or float(self.max_value) <= 0:
+                    raise ValueError(
+                        "min_value and max_value must be positive for log-uniform distribution"
+                    )
         elif self.param_type == ParamType.CATEGORICAL:
             if not self.choices:
                 raise ValueError("choices required for categorical parameters")
@@ -36,9 +42,16 @@ def integer(cls, min_val: int, max_val: int, default: int):
 
     @classmethod
     def float_range(cls, min_val: float, max_val: float, default: float):
-        """Create a float parameter space"""
+        """Create a float parameter space (uniform distribution)"""
         return cls(ParamType.FLOAT, min_val, max_val, default=default)
 
+    @classmethod
+    def float_log_range(cls, min_val: float, max_val: float, default: float):
+        """Create a float parameter space with log-uniform distribution."""
+        if min_val <= 0 or max_val <= 0:
+            raise ValueError("min_val and max_val must be positive for log-uniform distribution")
+        return cls(ParamType.FLOAT_LOG, min_val, max_val, default=default)
+
     @classmethod
     def categorical(cls, choices: List[Any], default: Any):
         """Create a categorical parameter space"""

models/cnn.py

@@ -8,6 +8,12 @@
 import torch
 import torch.nn as nn
 import torch.optim as optim
+from sklearn.metrics import (
+    accuracy_score,
+    classification_report,
+    f1_score,
+    roc_auc_score,
+)
 from torch.utils.data import DataLoader
 from torch.utils.tensorboard import SummaryWriter
 
@@ -34,7 +40,7 @@ class TrainingConfig:
     checkpoint_path: Path = MODEL_PATH
     grad_clip_norm: float = 1.0
     writer: Optional[SummaryWriter] = None
-    batch_size: int = 64
+    batch_size: int = 128
 
 
 class Backbone(nn.Module):
@@ -48,34 +54,39 @@ def __init__(
         stride: int = 1,
     ) -> None:
         super().__init__()
-
-        self.features = nn.Sequential(
+        padding = kernel_size // 2
+        self.block1 = nn.Sequential(
             nn.Conv2d(
-                in_channels, 32, kernel_size=kernel_size, stride=stride, padding=1
+                in_channels, 16, kernel_size=kernel_size, stride=stride, padding=padding
             ),
-            nn.BatchNorm2d(32),
+            nn.BatchNorm2d(16),
             nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        )
+
+        self.block2 = nn.Sequential(
             nn.Conv2d(
-                32, 64, kernel_size=kernel_size, stride=stride, padding=1
+                16, 32, kernel_size=kernel_size, stride=stride, padding=padding
             ),
-            nn.BatchNorm2d(64),
+            nn.BatchNorm2d(32),
             nn.ReLU(inplace=True),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        )
+
+        self.block3 = nn.Sequential(
             nn.Conv2d(
-                64,
-                128,
-                kernel_size=kernel_size,
-                stride=stride,
-                padding=1,
+                32, 64, kernel_size=kernel_size, stride=stride, padding=padding
             ),
-            nn.BatchNorm2d(128),
+            nn.BatchNorm2d(64),
             nn.ReLU(inplace=True),
             nn.AdaptiveAvgPool2d(1),
         )
 
+        self.features = nn.Sequential(self.block1, self.block2, self.block3)
+
         self.classifier = nn.Sequential(
             nn.Flatten(),
-            nn.Dropout(),
-            nn.Linear(128, num_classes),
+            nn.Linear(64, num_classes),
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -121,6 +132,7 @@ def train(
         grad_clip_norm: Optional[float] = None,
         writer: Optional[SummaryWriter] = None,
         num_workers: int = 2,
+        verbose: bool = True,
     ) -> Dict[str, float]:
         if self.network is None:
             raise RuntimeError("Train called before model is initialized")
@@ -166,8 +178,9 @@ def train(
         checkpoint = Checkpoint(str(config.checkpoint_path))
 
         total_params, trainable_params = count_parameters(self.network)
-        print(f"Total parameters: {total_params:,}")
-        print(f"Trainable parameters: {trainable_params:,}")
+        if verbose:
+            print(f"Total parameters: {total_params:,}")
+            print(f"Trainable parameters: {trainable_params:,}")
 
         history = {
             "train_loss": [],
@@ -177,7 +190,8 @@ def train(
         }
 
         for epoch in range(1, config.epochs + 1):
-            print(f"\nEpoch {epoch}/{config.epochs}")
+            if verbose:
+                print(f"\nEpoch {epoch}/{config.epochs}")
             train_loss, train_acc = train_epoch(
                 self.network,
                 train_loader,
@@ -188,6 +202,7 @@ def train(
                 epoch=epoch,
                 grad_clip_norm=config.grad_clip_norm,
                 writer=config.writer,
+                verbose=verbose,
             )
             val_loss, val_acc = validate(
                 self.network,
@@ -196,35 +211,40 @@ def train(
                 device,
                 epoch=epoch,
                 writer=config.writer,
+                verbose=verbose,
             )
 
             history["train_loss"].append(train_loss)
             history["train_acc"].append(train_acc)
             history["val_loss"].append(val_loss)
             history["val_acc"].append(val_acc)
 
-            print(
-                f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f} ({train_acc * 100:.2f}%)"
-            )
-            print(
-                f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f} ({val_acc * 100:.2f}%)"
-            )
+            if verbose:
+                print(
+                    f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f} ({train_acc * 100:.2f}%)"
+                )
+                print(
+                    f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f} ({val_acc * 100:.2f}%)"
+                )
 
             if checkpoint.save_if_better(
                 self.network, optimizer, epoch, val_acc, train_acc
             ):
-                print(
-                    f"Saved best model (val_acc={val_acc:.4f}) to {config.checkpoint_path}"
-                )
+                if verbose:
+                    print(
+                        f"Saved best model (val_acc={val_acc:.4f}) to {config.checkpoint_path}"
+                    )
 
             if early_stopper(val_loss, val_acc):
-                print(f"\nEarly stopping at epoch {epoch}")
+                if verbose:
+                    print(f"\nEarly stopping at epoch {epoch}")
                 break
 
-        print("\nTraining complete!")
-        print(
-            f"Best val acc: {checkpoint.best_val_acc:.4f} ({checkpoint.best_val_acc * 100:.2f}%)"
-        )
+        if verbose:
+            print("\nTraining complete!")
+            print(
+                f"Best val acc: {checkpoint.best_val_acc:.4f} ({checkpoint.best_val_acc * 100:.2f}%)"
+            )
 
         return {
             "best_val_acc": checkpoint.best_val_acc,
@@ -285,8 +305,9 @@ def evaluate(
         criterion = criterion or nn.CrossEntropyLoss()
 
         total_loss = 0.0
-        total_correct = 0
-        total_samples = 0
+        all_predictions = []
+        all_labels = []
+        all_probas = []
 
         with torch.no_grad():
             for images, labels in data_loader:
@@ -297,19 +318,45 @@ def evaluate(
 
                 total_loss += loss.item()
                 preds = torch.argmax(logits, dim=1)
-                total_correct += (preds == labels).sum().item()
-                total_samples += labels.size(0)
+                probas = torch.softmax(logits, dim=1)
+                
+                all_predictions.extend(preds.cpu().numpy())
+                all_labels.extend(labels.cpu().numpy())
+                all_probas.extend(probas.cpu().numpy())
+
+        y_true = np.array(all_labels)
+        y_pred = np.array(all_predictions)
+        y_proba = np.array(all_probas)
+
+        accuracy = accuracy_score(y_true, y_pred)
+        report = classification_report(
+            y_true, y_pred, output_dict=True, zero_division=0
+        )
+        
+        precision_macro = report["macro avg"]["precision"]
+        recall_macro = report["macro avg"]["recall"]
+        f1_macro = report["macro avg"]["f1-score"]
+        f1_micro = report.get("micro avg", {}).get("f1-score", f1_score(y_true, y_pred, average="micro", zero_division=0))
+        
+        roc_auc = roc_auc_score(y_true, y_proba, average="macro", multi_class="ovr")
 
         avg_loss = total_loss / len(data_loader)
-        accuracy = total_correct / total_samples if total_samples else 0.0
 
-        return {"loss": avg_loss, "accuracy": accuracy}
+        return {
+            "loss": avg_loss,
+            "accuracy": accuracy,
+            "precision_macro": precision_macro,
+            "recall_macro": recall_macro,
+            "f1_macro": f1_macro,
+            "f1_micro": f1_micro,
+            "roc_auc": roc_auc,
+        }
 
     def get_param_space(self) -> Dict[str, ParamSpace]:
         return {
             "kernel_size": ParamSpace.integer(min_val=3, max_val=5, default=3),
             "stride": ParamSpace.integer(min_val=1, max_val=3, default=1),
-            "learning_rate": ParamSpace.float_range(
+            "learning_rate": ParamSpace.float_log_range(
                 min_val=1e-5, max_val=1e-2, default=3e-4
             ),
             "batch_size": ParamSpace.categorical(choices=[16, 32, 64, 128], default=64),