fix: visualise batch & epoch metrics separately

trainingModel/Training.py

@@ -15,10 +15,10 @@ def train_model(
     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",
-    flatten_input = False,
-    num_classes : int = 39,
-
 ):
     """
     Trains the given model and returns:
@@ -40,7 +40,11 @@ def train_model(
 
     # Loss and optimizer
     criterion = nn.CrossEntropyLoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=lr ) # might add momentum 0.9 later
+
+    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)
@@ -49,20 +53,31 @@ def train_model(
     # 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)
+    #training_losses = np.zeros(num_batches * n_epochs)
+    #training_accuracies = np.zeros(num_batches * 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
 
+
     #----------------------
     # training loop
     #----------------------
@@ -71,8 +86,12 @@ def train_model(
         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 i, batch in enumerate(train_loader):
+        for batch in train_loader:
 
             # move to GPU memory
             inputs = batch["image"].to(device)
@@ -88,22 +107,30 @@ def train_model(
             outputs = model(inputs)
             loss = criterion(outputs, labels)
             
-            # Backward pass
+            # Backward pass & update params
             loss.backward()
-
-            # updates the parameters
             optimizer.step()
-            
-            # log the loss value
-            training_losses[epoch * num_batches + i] = loss.item() 
 
-            #updates the accuracy computation with new data 
-            train_accuracy_fn.update(outputs, labels)
+            # 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)
+
 
-            #compute accuracy with the current data
-            training_accuracies[epoch * num_batches + i] = train_accuracy_fn.compute().item() 
+        # Epoch-level metrics (average over all batches)
+        epoch_loss_avg = running_loss / running_total
+        epoch_acc_avg = running_correct / running_total
 
-        print(f'Epoch {epoch + 1} training complete')
+        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}')
 
         # ----------------------
         # validation loop
@@ -123,25 +150,30 @@ def train_model(
                     inputs = inputs.view(inputs.size(0), -1)
 
                 outputs = model(inputs)
-
                 val_accuracy_fn.update(outputs, labels)
 
-        current_accuracy = val_accuracy_fn.compute().item()
-        val_accuracies[epoch] = current_accuracy
+        
+        current_val_accuracy = val_accuracy_fn.compute().item()
+        val_accuracies[epoch] = current_val_accuracy
+
+        print(f"Epoch {epoch+1}: val acc={current_val_accuracy:.4f}")
 
         # keep track of best validation accuracy and save best model so far
-        if current_accuracy > best_accuracy:
-            best_accuracy = current_accuracy
+        if current_val_accuracy > best_accuracy:
+            best_accuracy = current_val_accuracy
             torch.save(model.state_dict(), save_path)
-            print(f'Epoch {epoch + 1} (validation accuracy: {best_accuracy})')
+        
+
         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,
+        "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,
     }

trainingModel/run_training.py

@@ -48,8 +48,6 @@ except Exception as e:
 
 full_dataset = ds['train']
 
-
-
 # Apply subset indices to full dataset - this gives you the same subset as data prep
 subset_dataset = full_dataset.select(subset_indices)
 
@@ -95,15 +93,24 @@ training_task = Task.init(
     reuse_last_task_id=False,
 )
 
+# Detail the data prep task used
 training_logger = training_task.get_logger()
-training_task.connect({"data_prep_task_used": DYNAMIC_TASK_ID})
+data_prep_metadata = {
+    "data_prep_task_id": DYNAMIC_TASK_ID,
+    "dataset_id": dataset_id,
+    "dataset_link": dataset_link,
+    "augmentation_used": aug_config,
+    "seed_used": seed,
+}
+training_task.connect(data_prep_metadata, name="data_prep_metadata")
 
 # Training parameters - Modify these to experiment
 training_config = {
     "num_classes": 39,
-    "n_epochs": 1,
+    "n_epochs": 3,
     "learning_rate": 1e-3,
     "batch_size": batch_size,
+    "optimizer": "adam",
     "save_path": "best_model.pt",
 }
 training_task.connect(training_config)
@@ -124,21 +131,30 @@ training_metrics = train_model(
     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"],
 )
 
 
 # ----------- Log metrics to ClearML -----------
 # Per-batch training losses and accuracies
-for i, loss in enumerate(training_metrics["losses"]):
-    training_logger.report_scalar("train", "loss_per_batch", value=loss, iteration=i)
+for i, loss in enumerate(training_metrics["batch_losses"]):
+    training_logger.report_scalar("train_batch", "loss", value=loss, iteration=i)
+
+for i, acc in enumerate(training_metrics["batch_accuracies"]):
+    training_logger.report_scalar("train_batch", "accuracy", value=acc, iteration=i)
+
 
-for i, acc in enumerate(training_metrics["accuracies"]):
-    training_logger.report_scalar("train", "accuracy_per_batch", 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("train_epoch", "loss", loss, iteration=epoch)
+    training_logger.report_scalar("train_epoch", "accuracy", acc, iteration=epoch)
 
-# Per-epoch validation accuracy
+# Per-epoch validation accuracies
 for epoch, acc in enumerate(training_metrics["val_accuracies"]):
-    training_logger.report_scalar("validation", "accuracy_per_epoch", value=acc, iteration=epoch)
+    training_logger.report_scalar("validation_epoch", "accuracy", value=acc, iteration=epoch)
 
 training_logger.report_single_value("best_val_accuracy", training_metrics["best_accuracy"])