Update train_test_utils/baseline.py

train_test_utils/baseline.py

@@ -1,10 +1,22 @@
 import torch
-import torch.nn as nn
-from torch.nn.utils import clip_grad_norm_
 from torch.utils.data import DataLoader
 
 
-def update_baseline(actor, baseline, validation_set, record_scores, batch_size=100, threshold=0.95):
+def update_baseline(actor, baseline, validation_set, record_scores=None, batch_size=100, threshold=0.95):
+    """
+    Evaluate the actor on the validation set and update the baseline if performance improves.
+    
+    Parameters:
+    - actor: current model being trained
+    - baseline: model used as the performance reference
+    - validation_set: dataset used for evaluation
+    - record_scores: previously recorded baseline scores
+    - batch_size: batch size for validation
+    - threshold: (optional) threshold for improvement (not used in current implementation)
+    
+    Returns:
+    - updated record_scores
+    """
 
     val_dataloader = DataLoader(dataset=validation_set,
                                 batch_size=batch_size,
@@ -17,25 +29,22 @@ def update_baseline(actor, baseline, validation_set, record_scores, batch_size=1
     for batch in val_dataloader:
         with torch.no_grad():
             actor_output = actor(batch)
-            actor_cost = actor_output['total_time']
-            actor_cost.reshape(-1)
-        actor_scores.append(actor_cost)
-    actor_scores = torch.cat(actor_scores, dim=0)
+            actor_cost = actor_output['total_time'].view(-1)
+            actor_scores.append(actor_cost)
 
+    actor_scores = torch.cat(actor_scores, dim=0)
+    actor_score_mean = actor_scores.mean().item()
 
     if record_scores is None:
         baseline.load_state_dict(actor.state_dict())
-        record_scores = actor_scores
-        return record_scores
-    else:
+        return actor_scores
 
-        if actor_scores.mean().item() < record_scores.mean().item():
-            print('\n', flush=True)
-            print('baseline updated', flush=True)
-            print('\n', flush=True)
+    baseline_score_mean = record_scores.mean().item()
 
-            baseline.load_state_dict(actor.state_dict())
-            record_scores = actor_scores
-            return record_scores
-        else:
-            return record_scores
+    if actor_score_mean < baseline_score_mean:
+        print(f"\nBaseline updated: {baseline_score_mean:.4f} → {actor_score_mean:.4f}\n", flush=True)
+        baseline.load_state_dict(actor.state_dict())
+        return actor_scores
+    else:
+        print(f"\nNo improvement: {actor_score_mean:.4f} ≥ {baseline_score_mean:.4f}\n", flush=True)
+        return record_scores