save always best model during validation

main_sphere_compression.py

@@ -327,20 +327,25 @@ def save_config(list_dicts, output_folder, filename=None):
                 f.write(f"{key:_<40}: {str(val)}\n")  # check this for all kinds of formatting
             f.write("="*60+"\n")
 
-def save_checkpoint(state_dics, is_best, output_folder, filename=None):
+def save_checkpoint(state_dics, is_best, output_folder, filename=None, only_best_model=False):
     os.makedirs(output_folder, exist_ok=True)
-    if filename:
-        if filename.endswith('.tar'): filename = filename.replace('.tar', '')
-        if filename.endswith('.pth'): filename = filename.replace('.pth', '')
-    else:
-        filename = f"checkpoint_{state_dics['epoch']+1:03d}"
-    output_fileAddr = os.path.join(output_folder, filename+'.pth.tar')
-    torch.save(state_dics, output_fileAddr)
-    print(f'saved checkpoint to {output_fileAddr}')
+    if not only_best_model:
+        if filename:
+            if filename.endswith('.tar'): filename = filename.replace('.tar', '')
+            if filename.endswith('.pth'): filename = filename.replace('.pth', '')
+        else:
+            filename = f"checkpoint_{state_dics['epoch']+1:04d}"
+        
+        output_fileAddr = os.path.join(output_folder, filename+'.pth.tar')
+        torch.save(state_dics, output_fileAddr)
+        print(f'saved checkpoint to {output_fileAddr}')
     if is_best:
         fp_best = os.path.join(output_folder, 'checkpoint_best_loss.pth.tar')
+        if only_best_model:
+            torch.save(state_dics, fp_best)
+        else:
+            shutil.copyfile(output_fileAddr, fp_best)
         print(f'saved best model to {fp_best}')
-        shutil.copyfile(output_fileAddr, fp_best)
 
 def get_model_name(string_name):
     for ch in [",", " ", "_", "-"]: # replace all with space
@@ -585,9 +590,10 @@ def main():
         if args.loss_file:
             with open(os.path.join(args.out_dir, args.loss_file), 'a') as f:
                 f.write(loss_str + '\n')
-
+        
+        perform_validation = args.validation_data and ((epoch + 1) >= args.validation_start)
         if args.validation_data:  # validation data
-            if (epoch + 1) >= args.validation_start:
+            if perform_validation:
                 saveVis = (((epoch+1) % args.interval_save_valtest == 0) or ((epoch+1) == args.max_epochs)) and (args.foldername_valtest is not None)
                 valTestVisFolder = os.path.join(args.out_dir, args.foldername_valtest) if saveVis else None
                 loss_validation = test_epoch(validation_dataloader, struct_loader, net, criterion, args.patch_res_valtest, valTestVisFolder, args.print_freq, epoch)
@@ -604,10 +610,11 @@ def main():
             if not args.no_scheduler:
                 scheduler.step() # type: ignore
                 scheduler_aux.step() # type: ignore
-
-        if (epoch+1) % args.checkpoint_interval == 0:
-            is_best = loss_validation["loss"] < best_loss if args.validation_data and ((epoch+1) >= args.validation_start) else False
-            best_loss = min(loss_validation["loss"], best_loss) if args.validation_data and ((epoch+1) >= args.validation_start) else best_loss
+        
+        save_regular_checkpoint = (epoch+1) % args.checkpoint_interval == 0
+        if save_regular_checkpoint or perform_validation:
+            is_best = loss_validation["loss"] < best_loss if perform_validation else False
+            best_loss = min(loss_validation["loss"], best_loss) if perform_validation else best_loss
             states = {
                 'epoch': epoch,
                 'net_state_dict': net.state_dict(),
@@ -619,7 +626,7 @@ def main():
             if not args.no_scheduler:
                 states["scheduler_state_dict"] = scheduler.state_dict()
                 states["scheduler_aux_state_dict"] = scheduler_aux.state_dict()
-            save_checkpoint(states, is_best, args.out_dir)
+            save_checkpoint(states, is_best, args.out_dir, only_best_model=(not save_regular_checkpoint))
 
     # Saving last results
     is_best = loss_validation["loss"] < best_loss if args.validation_data else False