Merge branch 'develop' of https://github.kcl.ac.uk/K23064919/smallGroupProject into develop

.gitignore

@@ -1,10 +1,8 @@
-<<<<<<< HEAD
 .vscode/
 .venv/
 .vscode/
 .models/
 __pycache__/
-=======
 
 # Python environment
 venv/
@@ -18,4 +16,4 @@ __pycache__/
 
 # Generated files from data_preparation.py
 class_distribution.png
->>>>>>> 04cb88662062ef6b880c627546d067fa0cedfa8b
+

best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23a4c08eaad4b40290eca84e6a8fa3e1d69bdf4312d5db6db5de96d1d8753024
+size 130261986

dataPrep/data_preparation.py

@@ -45,8 +45,9 @@ if torch.cuda.is_available():
 
 
 # ----- ClearML Setup -----
+project_name = "Small Group Project"
 task = Task.init(
-    project_name='Small Group Project',
+    project_name=f'{project_name}/Data Preparation',
     task_name='Data Preparation',
     task_type=Task.TaskTypes.data_processing
 )

dataPrep/helpers/clearml_data.py

@@ -11,12 +11,12 @@ Takes latest Data Prep ClearML task from project and reconstruct:
 - data loaders for both full and subset datasets
 - Aug settings used
 '''
-def extract_latest_data_task(project_name: str = "Small Group Project"):
+def extract_latest_data_task(project_name: str = "Small Group Project", num_workers: int = 8):
 
   # --------- Get latest Data Preparation task from ClearML ---------
 
   all_tasks = Task.get_tasks(
-    project_name=project_name,
+    project_name=f'{project_name}/Data Preparation',
     allow_archived=False,
     task_filter={'order_by': ["-last_update"]},
   )
@@ -76,7 +76,7 @@ def extract_latest_data_task(project_name: str = "Small Group Project"):
   subset_dataset = full_dataset.select(subset_indices)
 
   # Get data loaders for both full and subset datasets
-  subset_loaders, full_loaders, aug_config = get_data_loaders(data_params, subset_dataset, full_dataset)
+  subset_loaders, full_loaders, aug_config = get_data_loaders(data_params, subset_dataset, full_dataset, num_workers=num_workers)
   batch_size = int(data_params['General/dataloaders/batch_size'])
   seed = int(data_params['General/seed'])
 
@@ -99,7 +99,7 @@ def extract_latest_data_task(project_name: str = "Small Group Project"):
 Takes a given dataset, subset, data params to create DataLoaders
 Loaders split data into train, val, test
 '''
-def get_data_loaders(data_params, subset_dataset, full_dataset):
+def get_data_loaders(data_params, subset_dataset, full_dataset, num_workers):
 
   # Extract data parameters- these will be used in the DataLoaders
   seed = int(data_params['General/seed'])
@@ -115,7 +115,7 @@ def get_data_loaders(data_params, subset_dataset, full_dataset):
 
   # Create DataLoaders using the parameters from data prep
   subset_loaders = make_dataset_loaders(
-      subset_dataset, seed, batch_size, test_size, aug_config
+      subset_dataset, seed, batch_size, test_size, aug_config, workers=num_workers
   )
 
   print("\n--- Handoff Test Successful ---")
@@ -125,7 +125,7 @@ def get_data_loaders(data_params, subset_dataset, full_dataset):
 
 
   full_loaders = make_dataset_loaders(
-      full_dataset, seed, batch_size, test_size, aug_config
+      full_dataset, seed, batch_size, test_size, aug_config, workers=num_workers
   )
 
   print("\n--- Handoff Test Successful ---")

dataPrep/helpers/transforms_loaders.py

@@ -47,24 +47,25 @@ def make_augment_pipeline(aug_config):
     return augmentation
 
 
+def apply_augmentation(batch, augmentation):
+    batch['image'] = [augmentation(x) for x in batch['image']]
+    return batch
+
+def apply_normalisation(batch, normalisation):
+    batch['image'] = [normalisation(x) for x in batch['image']]
+    return batch
+    
+
 """
 Creates and returns DataLoaders (train, val, test) for a given dataset.
 Performs a 70/15/15 split
 """
-def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config):
+def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config, workers=8):
 
     # Define transformation pipelines for the dataset
     normalisation = make_norm_pipeline()
     augmentation = make_augment_pipeline(aug_config)
 
-    def apply_augmentation(batch):
-        batch['image'] = [augmentation(x) for x in batch['image']]
-        return batch
-
-    def apply_normalisation(batch):
-        batch['image'] = [normalisation(x) for x in batch['image']]
-        return batch
-    
     # 70/30 split creates train set
     split_1 = dataset.train_test_split(test_size=test_size, seed=seed)
     train_split = split_1['train']
@@ -76,14 +77,34 @@ def make_dataset_loaders(dataset, seed, batch_size, test_size, aug_config):
     val_split, test_split = split_2['train'], split_2['test']
 
     # Put each split through pipelines
-    train_split.set_transform(apply_augmentation)
-    val_split.set_transform(apply_normalisation)
-    test_split.set_transform(apply_normalisation)
+    train_split.set_transform(lambda batch: apply_augmentation(batch, augmentation))
+    val_split.set_transform(lambda batch: apply_normalisation(batch, normalisation))
+    test_split.set_transform(lambda batch: apply_normalisation(batch, normalisation))
 
     # Create dataloader for each
-    train_loader = DataLoader(train_split, batch_size=batch_size, shuffle=True)
-    val_loader = DataLoader(val_split, batch_size=batch_size, shuffle=False)
-    test_loader = DataLoader(test_split, batch_size=batch_size, shuffle=False)
+    train_loader = DataLoader(
+        train_split,
+        batch_size=batch_size,
+        shuffle=True,
+        pin_memory=True,
+        num_workers=workers
+    )
+    val_loader = DataLoader(
+        val_split,
+        batch_size=batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=workers
+    )
+    test_loader = DataLoader(
+        test_split,
+        batch_size=batch_size,
+        shuffle=False,
+        pin_memory=True,
+        num_workers=workers
+    )
+
+    print(f"\nWorkers used in DataLoaders: {workers}\n")
 
     dataset_loaders = {
         "train": train_loader,

models/modelTwo.py

@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class BetterCNN(nn.Module):
+    def __init__(self, noOfClasses=39):
+        super(BetterCNN, self).__init__()
+
+        # 32 Channels 
+        # We use padding=1 to keep spatial size same before pooling
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
+        self.bn1 = nn.BatchNorm2d(32)
+
+        # 64 Channels
+        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
+        self.bn2 = nn.BatchNorm2d(64)
+
+        # 128 Channels      
+        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
+        self.bn3 = nn.BatchNorm2d(128)
+
+        # 256 Channels
+        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
+        self.bn4 = nn.BatchNorm2d(256)
+
+        # Pooling layer 
+        self.pool = nn.MaxPool2d(2, 2)
+
+        # Adaptive Pooling  
+        self.adaptive_pool = nn.AdaptiveAvgPool2d((4, 4))
+
+        # Classification Head
+        self.fc1 = nn.Linear(256 * 4 * 4, 1024)
+        self.dropout = nn.Dropout(0.5) # Dropout after Linear layer
+        
+        self.fc2 = nn.Linear(1024, 512)
+        self.fc3 = nn.Linear(512, noOfClasses)
+
+    def forward(self, x):
+        # Block 1
+        x = self.conv1(x)
+        x = self.bn1(x)     # BatchNorm 
+        x = F.relu(x)
+        x = self.pool(x)    
+
+        # Block 2
+        x = self.pool(F.relu(self.bn2(self.conv2(x))))
+
+        # Block 3
+        x = self.pool(F.relu(self.bn3(self.conv3(x))))
+
+        # Block 4
+        x = self.pool(F.relu(self.bn4(self.conv4(x))))
+        
+        # Adapt & Flatten
+        x = self.adaptive_pool(x) 
+        x = torch.flatten(x, 1)   # Flattens to (Batch, 4096)
+        
+        # Dense Layers
+        x = F.relu(self.fc1(x))
+        x = self.dropout(x)       # Regularization
+        x = F.relu(self.fc2(x))
+        x = self.fc3(x)           # No activation needed here (handled by CrossEntropyLoss)
+
+        return x

subset_indices.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:972615a5b506b5ee2490f61866c26a4a2f9e2498c0baedb195a2a0d10a62e76f
+size 111016

testingModel/helpers/evaluation.py

@@ -0,0 +1,43 @@
+import torch
+from torch.nn import CrossEntropyLoss
+
+
+"""
+Evaluates a trained model on a dataloader that returns batches like:
+    batch["image"] -> Tensor [B, 3, 256, 256]
+    batch["label"] -> Tensor [B]
+
+Returns dict:
+    { "accuracy": float, "loss": float }
+"""
+def make_predictions(model, dataloader, device):
+
+    model.eval()
+    criterion = CrossEntropyLoss()
+
+    total_loss = 0
+    total_correct = 0
+    total_samples = 0
+
+    with torch.no_grad():
+        for batch in dataloader:
+
+            # Move tensors to device
+            images = batch["image"].to(device)
+            labels = batch["label"].to(device).long()
+
+            # Forward pass
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+
+            total_loss += loss.item() * images.size(0)
+            total_correct += (outputs.argmax(dim=1) == labels).sum().item()
+            total_samples += labels.size(0)
+
+    accuracy = total_correct / total_samples
+    avg_loss = total_loss / total_samples
+
+    return {
+        "accuracy": accuracy,
+        "loss": avg_loss,
+    }

testingModel/run_testing.py

@@ -0,0 +1,76 @@
+from clearml import Task
+from dataPrep.helpers.clearml_data import extract_latest_data_task
+
+import torch
+from models.modelOne import modelOne
+from testingModel.helpers.evaluation import make_predictions
+
+
+# -------------- Load Data --------------
+project_name = "Small Group Project"
+subset_loaders, full_loaders, data_prep_metadata = extract_latest_data_task(project_name=project_name)
+
+
+# -------- ClearML Testing Task Setup --------
+testing_task = Task.init(
+    project_name=f"{project_name}/Model Testing",
+    task_name="Model Testing",
+    task_type=Task.TaskTypes.testing,
+    reuse_last_task_id=False,
+)
+
+# Reference the data prep task used
+testing_logger = testing_task.get_logger()
+testing_task.connect(data_prep_metadata, name="data_prep_metadata_READONLY")
+
+CLEARML_TRAINING_ID = "5bac154a885b4acbaa07d8588027bb27"
+
+# Testing parameters - Modify these when experimenting
+testing_config = {
+    "model_train_id": CLEARML_TRAINING_ID,
+    "num_classes": 39,
+    "model_path": "best_model.pt",
+}
+testing_task.connect(testing_config)
+
+# Load the model weights from ClearML training task
+training_task = Task.get_task(task_id=testing_config["model_train_id"])
+model_artifact = training_task.artifacts.get("best_model")
+model_path = model_artifact.get_local_copy()
+
+# Reference training metadata
+training_hyperparams = training_task.get_parameters_as_dict()
+testing_task.connect(training_hyperparams['General'], name="training_metadata_READONLY")
+
+
+# -------- Rebuild the ML model --------
+model = modelOne()
+state_dict = torch.load(model_path, map_location="cpu") # Load to CPU first
+model.load_state_dict(state_dict)
+model.eval() # set dropout & batch norm layers to eval mode
+
+# Move model to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+
+# -------------------- Test model on test set --------------------
+testing_logger.report_text("Starting evaluation on TEST SUBSET...\n")
+test_subset = subset_loaders['test']
+
+subset_results = make_predictions(model, test_subset, device)
+
+
+# Accuracy & Loss logging
+testing_logger.report_single_value(name="Test Subset Accuracy", value=subset_results["accuracy"])
+testing_logger.report_single_value(name="Test Subset Loss", value=subset_results["loss"])
+
+
+# --------- Complete -----------------
+print("\n------ Testing Complete ------")
+testing_logger.report_text(
+    f"TEST SUBSET RESULTS:\n"
+    f"Loss: {subset_results['loss']:.4f}\n"
+    f"Accuracy: {subset_results['accuracy']:.4f}\n"
+)
+testing_task.close()

trainingModel/Training.py

@@ -1,182 +0,0 @@
-import torch 
-import torch.nn as nn
-import numpy as np
-from torcheval.metrics import MulticlassAccuracy
-from torch.utils.data import DataLoader
-
-
-# fix errors in runtime
-
-
-def train_model(
-    model: nn.Module,
-    train_loader: DataLoader,
-    val_loader: DataLoader,
-    device: torch.device,
-    n_epochs: int = 4,
-    lr: float = 1e-3,
-    num_classes: int = 39,
-    optimizer_type: str = "adam",
-    flatten_input: bool = False,
-    save_path: str = "best_model.pt",
-):
-    """
-    Trains the given model and returns:
-    - training_losses: numpy array of loss per batch
-    - training_accuracies: numpy array of running accuracy per batch
-    - val_accuracies: numpy array of accuracy per epoch
-    - best_accuracy: highest validation accuracy achieved
-
-    Expected batch format:
-        batch["image"] → Tensor [B, C, H, W]
-        batch["label"] → Tensor [B] with class IDs (int64)
-    Model output:
-        outputs → Tensor [B, num_classes] (logits)
-    """
-
-    
-    # Move model to device
-    model.to(device)
-
-    # Loss and optimizer
-    criterion = nn.CrossEntropyLoss()
-
-    if optimizer_type.lower() == "adam":
-        optimizer = torch.optim.Adam(model.parameters(), lr=lr ) # might add momentum 0.9 later
-    else:
-        optimizer = torch.optim.AdamW(model.parameters(), lr=lr )
-
-    # Metric trackers 
-    train_accuracy_fn = MulticlassAccuracy(num_classes=num_classes)
-    val_accuracy_fn = MulticlassAccuracy(num_classes=num_classes)
-
-    # Arrays to log metrics
-    num_batches = len(train_loader)
-
-    # Batch-level logs
-    batch_losses = []
-    batch_accuracies = []
-
-    # Epoch-level logs
-    epoch_losses = np.zeros(n_epochs)
-    epoch_accuracies = np.zeros(n_epochs)
-    val_accuracies = np.zeros(n_epochs)
-
-
-    if num_batches == 0:
-        raise RuntimeError("UH OH!!!! empty train loader")
-
-    # Store training losses and accuracies for every batch
-    # num_batches is the number of batches for every epoch
-    #training_losses = np.zeros(num_batches * n_epochs)
-    #training_accuracies = np.zeros(num_batches * n_epochs)
-
-    # store validation accuracy for every epoch
-
-
-    # keep track of best validation accuracy and best model
-    best_accuracy = 0.0
-
-
-    #----------------------
-    # training loop
-    #----------------------
-    
-    for epoch in range(n_epochs):
-        model.train()
-        train_accuracy_fn.reset()
-
-        running_loss = 0.0
-        running_correct = 0
-        running_total = 0
-
-        # iterate over all the dataloader's mini-batches
-        for batch in train_loader:
-
-            # move to GPU memory
-            inputs = batch["image"].to(device)
-            labels = batch["label"].to(device).long()
-
-            # flatten if not cnn REVISE LATER
-            if flatten_input:
-                inputs = inputs.view(inputs.size(0), -1)
-            
-            optimizer.zero_grad()
-
-            # Forward pass
-            outputs = model(inputs)
-            loss = criterion(outputs, labels)
-            
-            # Backward pass & update params
-            loss.backward()
-            optimizer.step()
-
-            # Log batch-level metrics
-            batch_losses.append(loss.item())
-            batch_acc = (outputs.argmax(dim=1) == labels).float().mean().item()
-            batch_accuracies.append(batch_acc)
-
-            # Sum epoch stats
-            running_loss += loss.item() * inputs.size(0)
-            running_correct += (outputs.argmax(dim=1) == labels).sum().item()
-            running_total += labels.size(0)
-
-
-        # Epoch-level metrics (average over all batches)
-        epoch_loss_avg = running_loss / running_total
-        epoch_acc_avg = running_correct / running_total
-
-        epoch_losses[epoch] = epoch_loss_avg
-        epoch_accuracies[epoch] = epoch_acc_avg
-
-        print(f"\n--- Epoch {epoch + 1}: ---")
-        print(f'Train loss={epoch_loss_avg:.4f}\nTrain accuracy={epoch_acc_avg:.4f}\n')
-
-        # ----------------------
-        # validation loop
-        # ----------------------
-
-        model.eval()
-        val_accuracy_fn.reset()
-
-
-        with torch.no_grad():
-            for batch in val_loader:
-                inputs = batch["image"].to(device)
-                labels = batch["label"].to(device).long()
-
-                # flatten if not cnn REVISE LATER
-                if flatten_input:
-                    inputs = inputs.view(inputs.size(0), -1)
-
-                outputs = model(inputs)
-                val_accuracy_fn.update(outputs, labels)
-
-        
-        current_val_accuracy = val_accuracy_fn.compute().item()
-        val_accuracies[epoch] = current_val_accuracy
-
-        print(f"\nEpoch {epoch+1}: val acc={current_val_accuracy:.4f}")
-
-        # keep track of best validation accuracy and save best model so far
-        if current_val_accuracy > best_accuracy:
-            best_accuracy = current_val_accuracy
-            torch.save(model.state_dict(), save_path)
-        
-
-        print(f'Epoch {epoch + 1} validation complete\n')
-
-    print(f"\nTraining finished. Best val accuracy: {best_accuracy:.4f}")
-    print(f"Best model weights saved to: {save_path}")
-
-    training_metrics = {
-        "batch_losses": np.array(batch_losses),
-        "batch_accuracies": np.array(batch_accuracies),
-        "epoch_losses": epoch_losses,
-        "epoch_accuracies": epoch_accuracies,
-        "val_accuracies": val_accuracies,
-        "best_accuracy": best_accuracy,
-    }
-
-    return training_metrics
-

trainingModel/helpers/Training.py

@@ -0,0 +1,199 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from torcheval.metrics import MulticlassAccuracy
+from torch.utils.data import DataLoader
+
+
+
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", DEVICE)
+
+def train_model(
+   model: nn.Module,
+   train_loader: DataLoader,
+   val_loader: DataLoader,
+   n_epochs: int = 4,
+   lr: float = 1e-3,
+   save_path: str = "best_model.pt",
+   num_classes : int = 39,
+   early_stop : int = 3,
+
+
+):
+   """
+   Trains the given model and returns:
+   - training_losses: numpy array of loss per epoch
+   - training_accuracies: numpy array of running accuracy per epoch
+   - val_accuracies: numpy array of accuracy per epoch
+   - best_accuracy: highest validation accuracy achieved
+
+
+   Expected batch format:
+       batch["image"] → Tensor [B, C, H, W]
+       batch["label"] → Tensor [B] with class IDs (int64)
+   Model output:
+       outputs → Tensor [B, num_classes] (logits)
+   """
+
+
+   # Move model to device
+   model.to(DEVICE)
+
+
+   # Loss and optimizer
+   criterion = nn.CrossEntropyLoss()
+   optimizer = torch.optim.Adam(model.parameters(), lr=lr ) # might add momentum 0.9 later
+
+
+   # Metric trackers
+   train_accuracy_fn = MulticlassAccuracy(num_classes=num_classes)
+   val_accuracy_fn = MulticlassAccuracy(num_classes=num_classes)
+
+
+   # Arrays to log metrics
+   num_batches = len(train_loader)
+
+
+   if num_batches == 0:
+       raise RuntimeError("UH OH!!!! empty train loader")
+
+
+   # Store training losses and accuracies for every epoch
+   training_losses = np.zeros(n_epochs)
+   training_accuracies = np.zeros(n_epochs)
+
+
+   # store validation accuracy for every epoch
+   val_accuracies = np.zeros(n_epochs)
+
+
+   # keep track of best validation accuracy and best model
+   best_accuracy = 0.0
+
+
+   # keep track of accuracy improvement
+   improv_counter = 0
+
+
+   #----------------------
+   # training loop
+   #----------------------
+  
+   for epoch in range(n_epochs):
+       model.train()
+       train_accuracy_fn.reset()
+
+
+       training_loss = 0.0
+
+
+       # iterate over all the dataloader's mini-batches
+       for i, batch in enumerate(train_loader):
+
+
+           # move to GPU memory
+           inputs = batch["image"].to(DEVICE)
+           labels = batch["label"].to(DEVICE).long()
+
+
+
+
+           optimizer.zero_grad()
+
+
+           # Forward pass
+           outputs = model(inputs)
+           loss = criterion(outputs, labels)
+          
+           # Backward pass
+           loss.backward()
+
+
+           # updates the parameters
+           optimizer.step()
+          
+           # log the loss value for epoch
+           training_loss += loss.item()
+
+
+           #updates the accuracy computation with new data
+           train_accuracy_fn.update(outputs, labels)
+
+
+       # compute epoch-level training metrics
+       training_losses[epoch] = training_loss / num_batches
+       training_accuracies[epoch] = train_accuracy_fn.compute().item()
+
+
+       print(f'Epoch {epoch + 1} training complete. Training Accuracy: {training_accuracies[epoch]:.4f}')
+
+
+       # ----------------------
+       # validation loop
+       # ----------------------
+
+
+       model.eval()
+       val_accuracy_fn.reset()
+
+
+
+
+       with torch.no_grad():
+           for batch in val_loader:
+               inputs = batch["image"].to(DEVICE)
+               labels = batch["label"].to(DEVICE).long()
+
+
+               outputs = model(inputs)
+
+
+               val_accuracy_fn.update(outputs, labels)
+
+
+       current_accuracy = val_accuracy_fn.compute().item()
+       val_accuracies[epoch] = current_accuracy
+
+
+       # keep track of best validation accuracy and save best model so far
+       if current_accuracy > best_accuracy:
+           best_accuracy = current_accuracy
+           torch.save(model.state_dict(), save_path)
+           improv_counter = 0  #Resets coounter if accuracy improves
+           print(f'Epoch {epoch + 1} (validation accuracy: {best_accuracy})')
+
+
+       else:
+           improv_counter +=1
+           print(f'No improvement for {improv_counter} epoch')
+
+
+           if improv_counter >= early_stop:
+               print (f"Early stopping at epoch {epoch +1}")
+               break
+
+
+
+
+       print(f'Epoch {epoch + 1} validation complete')
+
+
+   print(f"\nTraining finished. Best val accuracy: {best_accuracy:.4f}")
+   print(f"Best model weights saved to: {save_path}")
+
+   
+   training_metrics = {
+       "losses": training_losses,
+       "accuracies": training_accuracies,
+       "val_accuracies": val_accuracies,
+       "best_accuracy": best_accuracy
+
+   }
+
+   return training_metrics
+
+
+
+

trainingModel/run_training.py

@@ -1,20 +1,21 @@
-
+import os
 from clearml import Task
 from dataPrep.helpers.clearml_data import extract_latest_data_task
 
 import torch
-from models.modelOne import modelOne
-from trainingModel.Training import train_model
+from models.modelTwo import BetterCNN
+from trainingModel.helpers.Training import train_model
 
 
 # -------------- Load Data --------------
+NUM_WORKERS = 0
 project_name = "Small Group Project"
-subset_loaders, full_loaders, data_prep_metadata = extract_latest_data_task(project_name=project_name)
+subset_loaders, full_loaders, data_prep_metadata = extract_latest_data_task(project_name=project_name, num_workers=NUM_WORKERS)
 
 
 # -------- ClearML Training Task Setup --------
 training_task = Task.init(
-    project_name="Small Group Project",
+    project_name=f"{project_name}/Model Training",
     task_name="Model Training",
     reuse_last_task_id=False,
 )
@@ -26,18 +27,24 @@ training_task.connect(data_prep_metadata, name="data_prep_metadata_READONLY")
 # Training parameters - Modify these to experiment
 training_config = {
     "num_classes": 39,
-    "n_epochs": 10,
+    "n_epochs": 1,
     "learning_rate": 1e-3,
     "optimizer": "adam",
     "save_path": "best_model.pt",
+    "num_workers": NUM_WORKERS
 }
 training_task.connect(training_config)
 
 
 # -------- Build the ML model --------
-model = modelOne(noOfClasses=training_config["num_classes"])
+model = BetterCNN(noOfClasses=training_config["num_classes"])
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
+# Print device info
+print(f"\n**Using device: {device}**\n") 
+if device.type == 'cuda':
+    print(f"GPU Name: {torch.cuda.get_device_name(0)}")
 
 # ------- Train the model (on subset for now) -------
  
@@ -46,33 +53,27 @@ training_metrics = train_model(
     model=model,
     train_loader=subset_loaders['train'],
     val_loader=subset_loaders['val'],
-    device=device,
     n_epochs=training_config["n_epochs"],
     lr=training_config["learning_rate"],
     num_classes=training_config["num_classes"],
-    optimizer_type=training_config["optimizer"],
     save_path=training_config["save_path"],
+    early_stop=3,
 )
 
 
 # ----------- Log metrics to ClearML -----------
-# Per-batch training losses and accuracies
-for i, loss in enumerate(training_metrics["batch_losses"]):
-    training_logger.report_scalar("training batch loss", "loss", value=loss, iteration=i)
-
-for i, acc in enumerate(training_metrics["batch_accuracies"]):
-    training_logger.report_scalar("training batch accuracy", "accuracy", value=acc, iteration=i)
-
 # Per-epoch training losses and accuracies
-epoch_metrics = zip(training_metrics["epoch_losses"], training_metrics["epoch_accuracies"])
-for epoch, (loss, acc) in enumerate(epoch_metrics):
-    training_logger.report_scalar("training epoch loss", "loss", loss, iteration=epoch)
-    training_logger.report_scalar("training epoch accuracy", "accuracy", acc, iteration=epoch)
+for epoch, loss in enumerate(training_metrics["losses"]):
+    training_logger.report_scalar("training epoch loss", "loss", value=loss, iteration=epoch)
+
+for epoch, acc in enumerate(training_metrics["accuracies"]):
+    training_logger.report_scalar("training epoch accuracy", "accuracy", value=acc, iteration=epoch)
 
 # Per-epoch validation accuracies
 for epoch, acc in enumerate(training_metrics["val_accuracies"]):
     training_logger.report_scalar("validation epoch accuracy", "accuracy", value=acc, iteration=epoch)
 
+# Best validation accuracy
 training_logger.report_single_value("best_val_accuracy", training_metrics["best_accuracy"])
 
 # Upload best model as artifact