Fix best-checkpoint saving bug + evaluation fallback + cleanup

train_and_test.py

@@ -6,6 +6,7 @@ For further details onto which params file to load
 read the README in `params_storage` folder.
 '''
 import sys
+import os
 import torch
 import warnings
 with warnings.catch_warnings():
@@ -235,6 +236,15 @@ Model Evaluation at the best epoch
 epoch = "best"  # int/"best"/"last"
 eval_name = opt_exp.checkpoints_dir
 
+# Check if "best" checkpoint exists; fall back to "latest" if not
+if model_type == 'mamba':
+    best_path = os.path.join(eval_name, 'mamba_dloc', f'best_net_mamba_dloc.pth')
+else:
+    best_path = os.path.join(eval_name, 'encoder', f'best_net_encoder.pth')
+if not os.path.exists(best_path):
+    print(f'WARNING: "best" checkpoint not found at {best_path}, falling back to "latest"')
+    epoch = "latest"
+
 if model_type == 'mamba':
     # Mamba: reload single model
     mamba_model.load_networks(epoch, load_dir=eval_name)

trainer.py

@@ -32,7 +32,7 @@ def train(model, train_loader, test_loader):
     # initialization
     total_steps = 0
     print('Training called')
-    stopping_count = 0
+    best_median = float('inf')
 
     for epoch in range(model.opt.starting_epoch_count+1, model.opt.n_epochs+1): # opt.niter + opt.niter_decay + 1):
         epoch_start_time = time.time()
@@ -79,21 +79,21 @@ def train(model, train_loader, test_loader):
         write_log([str(nighty_percentile_error_tr)], model.decoder.model_name, log_dir=model.decoder.opt.log_dir, log_type='train_90_error')
         if (epoch == model.opt.starting_epoch_count + 1):
             min_eval_loss, median_error = test(model, test_loader, save_output=False)
+            best_median = median_error
+            print('Saving initial best model at epoch %d (median=%.4f)' % (epoch, best_median))
+            model.save_networks('best')
         else:
             new_eval_loss, new_med_error = test(model, test_loader, save_output=False)
-            if (median_error>=new_med_error):
-                stopping_count = stopping_count+1
-                median_error = new_med_error
+            if new_med_error < best_median:
+                best_median = new_med_error
+                print('Saving new best model at epoch %d (median=%.4f)' % (epoch, best_median))
+                model.save_networks('best')
+            median_error = new_med_error
 
-        # generated_outputs = temp_generator_outputs
         if epoch % model.encoder.opt.save_epoch_freq == 0:
             print('saving the model at the end of epoch %d, iters %d' %(epoch, total_steps))
             model.save_networks('latest')
             model.save_networks(epoch)
-            if (stopping_count==2):
-                print('Saving best model at %d epoch' %(epoch))
-                model.save_networks('best')
-                stopping_count=0
 
         print('End of epoch %d / %d \t Time Taken: %d sec' %
               (epoch, model.decoder.opt.niter + model.decoder.opt.niter_decay, time.time() - epoch_start_time))