Update Training.py

trainingModel/Training.py

@@ -2,16 +2,10 @@ import torch
 import torch.nn as nn
 import numpy as np
 from torcheval.metrics import MulticlassAccuracy
-#from torchvision import transforms
-
-
-
 from torch.utils.data import DataLoader
-#from torchvision.datasets import MNIST
 
-#import torchvision.utils
 
-# loss, optimizer, training loop, validation, best model saving
+# fix errors in runtime
 
 
 def train_model(
@@ -26,7 +20,19 @@ def train_model(
     num_classes : int = 39,
 
 ):
-
+    """
+    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
@@ -43,19 +49,20 @@ def train_model(
     # 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 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
     val_accuracies = np.zeros(n_epochs)
+
     # keep track of best validation accuracy and best model
     best_accuracy = 0.0
 
-
-
     #----------------------
     # training loop
     #----------------------
@@ -69,16 +76,14 @@ def train_model(
 
             # move to GPU memory
             inputs = batch["image"].to(device)
-            labels = batch["label"].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)
@@ -92,40 +97,31 @@ def train_model(
             # log the loss value
             training_losses[epoch * num_batches + i] = loss.item() 
 
-            # Compute accuracy of the batch.
-            
-
             #updates the accuracy computation with new data 
             train_accuracy_fn.update(outputs, labels)
 
             #compute accuracy with the current data
             training_accuracies[epoch * num_batches + i] = train_accuracy_fn.compute().item() 
 
-            
-            # display some progress (every 200 batches)
-            # optional, you can comment out
-            # if i % 200 == 0:
-            #     print(f'Epoch {epoch + 1}, batch {i+1} of {len(train_loader)}')
-
         print(f'Epoch {epoch + 1} training complete')
 
-        # Validation after each epoch
+        # ----------------------
+        # validation loop
+        # ----------------------
+
         model.eval()
         val_accuracy_fn.reset()
 
 
-        # The context 'torch.no_grad()' tells pytorch we are not interested in computing 
-        # gradients here, so forward pass is more efficient
         with torch.no_grad():
-            for i, batch in enumerate(val_loader):
+            for batch in val_loader:
                 inputs = batch["image"].to(device)
-                labels = batch["label"].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)
@@ -133,7 +129,6 @@ def train_model(
         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
@@ -146,8 +141,3 @@ def train_model(
 
     return training_losses, training_accuracies, val_accuracies, best_accuracy
 
-
-    #tweak later
-    #best_model = MNISTNet().to(device)
-    #best_model.load_state_dict(
-    #    torch.load('mnist-torch-best_model.pt', map_location=device))