Merge pull request #2 from AI-Enabled-Software-Testing/feat/ml-models

Update models

Author: seofernando25

.gitignore

@@ -3,4 +3,5 @@ __pycache__/
 .cache/
 .ruff_cache/
 *.pyc
-.aim/
+.aim/
+**/runs/

README.md

@@ -5,11 +5,6 @@ This project aims to explore and analyze metaheuristic search-based algorithms f
 This is our [idea](./Project%20Proposal/Project%20Proposal%20-%20Fernando%20and%20Kelvin.pdf).
 
 ## Datasets
-* [MNIST](https://scikit-learn.org/stable/modules/generated/sklearn.datasets.fetch_openml.html)
-    * Handwritten Digit Recognition
-    * Using scikit-learn's fetch_openml
-    * 28x28 Grayscale Images
-    * 10 Classes of digits (0-9)
 * [CIFAR-10](https://www.cs.toronto.edu/~kriz/cifar.html)
     * Object Recognition
     * 32x32 Colored Images
@@ -18,8 +13,7 @@ This is our [idea](./Project%20Proposal/Project%20Proposal%20-%20Fernando%20and%
 
 ## Models in Consideration
 * **Tree-based Model**: Decision Tree
-* **Linear/Polynomial-based**: Linear Regression (optional)
-* **Permutation-based** (especially, neural networks): Multi-Layer Perceptron (MLP)
+* **(Pixels) Permutation-based** (especially, neural networks): Convolutional Neural Network (CNN)
 * **Kernel-based**: K-Nearest-Neighbor (KNN)
 
 ## Metaheuristic Guided Search
@@ -81,4 +75,33 @@ This is our [idea](./Project%20Proposal/Project%20Proposal%20-%20Fernando%20and%
 2. Run `data_process.py` to process the images in the datasets. 
 3. Run `data_explorer.py` to view details of processed images from different API endpoints. 
    * Note: You may need to use a client such as Postman to launch those API requests. 
-   * Note: Refer to [`openapi.yaml`](openapi.yaml) for more detailed descriptions of those endpoints. 
+   * Note: Refer to [`openapi.yaml`](openapi.yaml) for more detailed descriptions of those endpoints. 
+### Quick Model Training
+You can quickly train a CNN model on CIFAR-10 using the provided training script:
+
+```bash
+# Basic training with default parameters (300 epochs)
+python scripts/train_cnn.py
+
+# Quick training with fewer epochs for testing
+python scripts/train_cnn.py --epochs 10 --batch-size 64
+
+# Custom training with specific hyperparameters
+python scripts/train_cnn.py --epochs 50 --batch-size 128 --lr 0.001 --model-path .cache/models/my_cnn.pth
+```
+
+**Available arguments:**
+- `--epochs`: Number of training epochs (default: 300)
+- `--batch-size`: Batch size for training (default: 128)
+- `--lr`: Learning rate (default: 0.0003)
+- `--model-path`: Path to save the trained model (default: .cache/models/cnn_cifar.pth)
+- `--device`: Force device selection (cuda/cpu, auto-detects if not specified)
+
+The script includes:
+- Automatic CIFAR-10 data loading and preprocessing
+- TensorBoard logging for training visualization
+- Early stopping and model checkpointing
+- CUDA support with automatic device detection
+
+### Model Training with a Customized Tuning Process
+* A Proof-of-Concept end-to-end quick demo is shown in the Jupyter Notebook: `notebooks\model_training_flow.ipynb`, including: a shorter demo with less data, data and model loading processes, an exhaustive tuning (without metaheuristics) on only the validation set, training and evaluating on the best found set of hyperparameters for each model.  

framework/data_utils.py

@@ -1,12 +1,15 @@
 """Data loading and preprocessing utilities."""
+
 from pathlib import Path
-from typing import Tuple, List
+from typing import 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 framework.datasets import CIFAR10Dataset
+
 from framework import utils
+from framework.datasets import CIFAR10Dataset
+
 
 def load_cifar10_data():
     """Load CIFAR-10 dataset (grayscale from processed datasets)."""
@@ -21,8 +24,7 @@ def load_cifar10_data():
             "uv run python -m scripts.data_process"
         )
 
-    ds_dict = load_from_disk(str(dataset_path))
-    return ds_dict
+    return load_from_disk(str(dataset_path))
 
 
 def prepare_data(ds_dict, split: str):
@@ -39,14 +41,10 @@ def split_train_val(
     images: List[np.ndarray],
     labels: np.ndarray,
     val_ratio: float = 0.2,
-    random_state: int = 42
+    random_state: int = 42,
 ) -> Tuple[List[np.ndarray], np.ndarray, List[np.ndarray], np.ndarray]:
     X_train, X_val, y_train, y_val = train_test_split(
-        images,
-        labels,
-        test_size=val_ratio,
-        stratify=labels,
-        random_state=random_state
+        images, labels, test_size=val_ratio, stratify=labels, random_state=random_state
     )
     return X_train, y_train, X_val, y_val
 
@@ -57,15 +55,15 @@ def create_dataloaders(
     X_val: List[np.ndarray],
     y_val: np.ndarray,
     batch_size: int,
-    num_workers: int = 2
+    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(),
+        pin_memory=utils.is_cuda_available(),
     )
 
     val_dataset = CIFAR10Dataset(X_val, y_val)
@@ -74,8 +72,7 @@ def create_dataloaders(
         batch_size=batch_size,
         shuffle=False,
         num_workers=num_workers,
-        pin_memory=utils.is_cuda()
+        pin_memory=utils.is_cuda_available(),
     )
 
     return train_loader, val_loader
-

framework/datasets.py

@@ -18,4 +18,3 @@ def __getitem__(self, idx):
         image = torch.from_numpy(image).unsqueeze(0)
 
         return image, label
-

framework/training.py

@@ -1,4 +1,5 @@
 import torch
+import torch.nn as nn
 from tqdm import tqdm
 from typing import Tuple, Optional
 from pathlib import Path
@@ -13,64 +14,56 @@ class Checkpoint:
     def __init__(self, model_path: str):
         self.model_path = Path(model_path)
         self.best_val_acc = 0.0
-    
+
     def save_if_better(
         self,
         model: Module,
         optimizer: Optimizer,
         epoch: int,
         val_acc: float,
         train_acc: float,
-        **kwargs
+        **kwargs,
     ) -> bool:
         """Save checkpoint if validation accuracy improved."""
         if val_acc > self.best_val_acc:
             self.best_val_acc = val_acc
             self.model_path.parent.mkdir(parents=True, exist_ok=True)
-            
-            checkpoint = {
-                'epoch': epoch,
-                'model_state_dict': model.state_dict(),
-                'optimizer_state_dict': optimizer.state_dict(),
-                'val_acc': val_acc,
-                'train_acc': train_acc,
-                **kwargs
+
+            checkpoint_data = {
+                "model_state_dict": model.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "epoch": epoch,
+                "val_acc": val_acc,
+                "train_acc": train_acc,
+                **kwargs,
             }
-            
-            torch.save(checkpoint, str(self.model_path))
+            torch.save(checkpoint_data, self.model_path)
             return True
         return False
-    
-    def load(self, model: Module, optimizer: Optimizer) -> dict:
-        """Load checkpoint from disk."""
-        checkpoint = torch.load(str(self.model_path), map_location='cpu')
-        model.load_state_dict(checkpoint['model_state_dict'])
-        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
-        self.best_val_acc = checkpoint.get('val_acc', 0.0)
-        return checkpoint
 
 
 class EarlyStopping:
-    def __init__(self, patience: int = 10, min_delta: float = 0.0) -> None:
+    def __init__(self, patience: int = 7, min_delta: float = 0.0):
         self.patience = patience
         self.min_delta = min_delta
-        self.best_loss = float('inf')
         self.counter = 0
-        self.best_acc = 0.0
-    
-    def __call__(self, val_loss: float, val_acc: float) -> bool:
-        if val_loss < self.best_loss - self.min_delta:
+        self.best_loss: Optional[float] = None
+        self.best_acc: float = 0.0
+
+    def __call__(self, val_loss: float, val_acc: Optional[float] = None) -> bool:
+        """Returns True if training should stop."""
+        if self.best_loss is None or val_loss < self.best_loss - self.min_delta:
             self.best_loss = val_loss
-            self.best_acc = val_acc
+            if val_acc is not None:
+                self.best_acc = max(self.best_acc, val_acc)
             self.counter = 0
             return False
-        else:
-            self.counter += 1
-            if self.counter >= self.patience:
-                return True
-            return False
 
+        if val_acc is not None:
+            self.best_acc = max(self.best_acc, val_acc)
 
+        self.counter += 1
+        return self.counter >= self.patience
 
 
 def train_epoch(
@@ -79,18 +72,20 @@ def train_epoch(
     criterion: torch.nn.Module,
     optimizer: Optimizer,
     device: torch.device,
-    scheduler: LRScheduler,
+    scheduler: Optional[LRScheduler] = None,
     epoch: int = 0,
     grad_clip_norm: float = 1.0,
-    writer: Optional[SummaryWriter] = None
+    writer: Optional[SummaryWriter] = None,
 ) -> Tuple[float, float]:
     """Trains the model for one epoch and returns the epoch loss and accuracy."""
-    model.train()
+    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)):
+    for batch_idx, (images, labels) in enumerate(
+        tqdm(train_loader, desc="Training", leave=False)
+    ):
         images, labels = images.to(device), labels.to(device)
 
         optimizer.zero_grad()
@@ -99,40 +94,32 @@ def train_epoch(
         loss.backward()
         torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip_norm)
         optimizer.step()
-        scheduler.step()
+        if scheduler is not None:
+            scheduler.step()
 
-        # Stats - compute once and reuse to avoid duplicate .item() calls
-        loss_value = loss.item()  # Single GPU->CPU sync
+        loss_value = loss.item()
         running_loss += loss_value
-        
+
         _, predicted = torch.max(outputs.data, 1)
-        batch_correct = (predicted == labels).sum().item()  # Single GPU->CPU sync
+        batch_correct = (predicted == labels).sum().item()
         total += labels.size(0)
         correct += batch_correct
 
-        # Track metrics
         if batch_idx % 10 == 0 and writer is not None:
             batch_total = labels.size(0)
             batch_acc = 100 * batch_correct / batch_total
-            current_lr = optimizer.param_groups[0]['lr']
+            current_lr = optimizer.param_groups[0]["lr"]
             step = epoch * len(train_loader) + batch_idx
-            writer.add_scalar('train/batch_loss', loss_value, step)
-            writer.add_scalar('train/batch_accuracy', batch_acc, step)
-            writer.add_scalar('train/learning_rate', current_lr, step)
+            writer.add_scalar("train/batch_loss", loss_value, step)
+            writer.add_scalar("train/batch_accuracy", batch_acc, step)
+            writer.add_scalar("train/learning_rate", current_lr, step)
 
     epoch_loss = running_loss / len(train_loader)
     epoch_acc = correct / total
 
-    # Track epoch-level metrics
     if writer is not None:
-        writer.add_scalar('train/epoch_loss', epoch_loss, epoch)
-        writer.add_scalar('train/epoch_accuracy', epoch_acc * 100, epoch)
-    
-    # Log parameter and gradient histograms (only every N epochs to reduce CPU overhead)
-    if writer is not None and (epoch % 10 == 0 or epoch == 1):  # Log every 10 epochs or first epoch
-        for name, param in model.named_parameters():
-            writer.add_histogram(f'train_params/{name}', param.data, epoch)
-            writer.add_histogram(f'train_grads/{name}', param.grad.data, epoch)
+        writer.add_scalar("train/epoch_loss", epoch_loss, epoch)
+        writer.add_scalar("train/epoch_accuracy", epoch_acc * 100, epoch)
 
     return epoch_loss, epoch_acc
 
@@ -143,7 +130,7 @@ def validate(
     criterion: torch.nn.Module,
     device: torch.device,
     epoch: int = 0,
-    writer: Optional[SummaryWriter] = None
+    writer: Optional[SummaryWriter] = None,
 ) -> Tuple[float, float]:
     """Validates the model and returns the epoch loss and accuracy."""
     model.eval()
@@ -157,8 +144,8 @@ def validate(
 
             outputs = model(images)
             loss = criterion(outputs, labels)
-
             running_loss += loss.item()
+
             _, predicted = torch.max(outputs.data, 1)
             total += labels.size(0)
             correct += (predicted == labels).sum().item()
@@ -167,8 +154,7 @@ def validate(
     epoch_acc = correct / total
 
     if writer is not None:
-        writer.add_scalar('val/epoch_loss', epoch_loss, epoch)
-        writer.add_scalar('val/epoch_accuracy', epoch_acc * 100, epoch)
+        writer.add_scalar("val/epoch_loss", epoch_loss, epoch)
+        writer.add_scalar("val/epoch_accuracy", epoch_acc * 100, epoch)
 
     return epoch_loss, epoch_acc
-