Merge pull request #8 from K23064919/test-model

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
 )

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/run_training.py

@@ -4,7 +4,7 @@ from dataPrep.helpers.clearml_data import extract_latest_data_task
 
 import torch
 from models.modelOne import modelOne
-from trainingModel.Training import train_model
+from trainingModel.helpers.Training import train_model
 
 
 # -------------- Load Data --------------
@@ -14,7 +14,7 @@ subset_loaders, full_loaders, data_prep_metadata = extract_latest_data_task(proj
 
 # -------- 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,7 +26,7 @@ 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",
@@ -37,7 +37,12 @@ training_task.connect(training_config)
 # -------- Build the ML model --------
 model = modelOne(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) -------