Update training loop and fixed training metrics

trainingModel/helpers/Training.py

@@ -1,182 +1,199 @@
-import torch 
+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
 
 
+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,
-    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",
+   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 batch
-    - training_accuracies: numpy array of running accuracy per batch
-    - val_accuracies: numpy array of accuracy per epoch
-    - best_accuracy: highest validation accuracy achieved
+   """
+   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()
 
-    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)
+           # updates the parameters
+           optimizer.step()
+          
+           # log the loss value for epoch
+           training_loss += loss.item()
 
-    # 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 )
+           #updates the accuracy computation with new data
+           train_accuracy_fn.update(outputs, labels)
 
-    # 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)
+       # compute epoch-level training metrics
+       training_losses[epoch] = training_loss / num_batches
+       training_accuracies[epoch] = train_accuracy_fn.compute().item()
 
-    # 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)
+       print(f'Epoch {epoch + 1} training complete. Training Accuracy: {training_accuracies[epoch]:.4f}')
 
 
-    if num_batches == 0:
-        raise RuntimeError("UH OH!!!! empty train loader")
+       # ----------------------
+       # validation loop
+       # ----------------------
 
-    # 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
+       model.eval()
+       val_accuracy_fn.reset()
 
 
-    # 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()
+       with torch.no_grad():
+           for batch in val_loader:
+               inputs = batch["image"].to(DEVICE)
+               labels = batch["label"].to(DEVICE).long()
 
-        running_loss = 0.0
-        running_correct = 0
-        running_total = 0
 
-        # iterate over all the dataloader's mini-batches
-        for batch in train_loader:
+               outputs = model(inputs)
 
-            # 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()
+               val_accuracy_fn.update(outputs, labels)
 
-            # 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)
+       current_accuracy = val_accuracy_fn.compute().item()
+       val_accuracies[epoch] = current_accuracy
 
-            # Sum epoch stats
-            running_loss += loss.item() * inputs.size(0)
-            running_correct += (outputs.argmax(dim=1) == labels).sum().item()
-            running_total += labels.size(0)
 
+       # 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})')
 
-        # 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
+       else:
+           improv_counter +=1
+           print(f'No improvement for {improv_counter} epoch')
 
-        print(f"\n--- Epoch {epoch + 1}: ---")
-        print(f'Train loss={epoch_loss_avg:.4f}\nTrain accuracy={epoch_acc_avg:.4f}\n')
 
-        # ----------------------
-        # validation loop
-        # ----------------------
+           if improv_counter >= early_stop:
+               print (f"Early stopping at epoch {epoch +1}")
+               break
 
-        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)
+       print(f'Epoch {epoch + 1} validation complete')
 
-                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"\nTraining finished. Best val accuracy: {best_accuracy:.4f}")
+   print(f"Best model weights saved to: {save_path}")
 
-        print(f"\nEpoch {epoch+1}: val acc={current_val_accuracy:.4f}")
+   
+   training_metrics = {
+       "losses": training_losses,
+       "accuracies": training_accuracies,
+       "val_accuracies": val_accuracies,
+       "best_accuracy": best_accuracy
 
-        # 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')
+   return training_metrics
 
-    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/run_training.py

@@ -53,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