Fix folder hierarchy for footsteps

uncbiag · Oct 29, 2024 · f2e870f · f2e870f
1 parent f78e27f
commit f2e870f
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/training/train_multi.py b/training/train_multi.py
@@ -323,14 +323,14 @@ def finetune(net, data_loader, val_data_loader, GPUS, epochs, eval_period, save_
 
     net_par.train()
 
-    footsteps.output_dir_impl = footsteps.output_dir + "finetune/"
+    footsteps.output_dir_impl = footsteps.output_dir.split("2nd_step/")[0] + "finetune/"
     os.makedirs(footsteps.output_dir + "checkpoints", exist_ok=True)
 
     train(net_par, optimizer, data_loader, val_data_loader, unwrapped_net=net, epochs=epochs, eval_period=eval_period, save_period=save_period, data_augmenter=augment)
 
     torch.save(
                 net.regis_net.state_dict(),
-                footsteps.output_dir + "checkpoints/Step_2_final.trch",
+                footsteps.output_dir + "checkpoints/finetune_final.trch",
             )
 
 if __name__ == "__main__":
@@ -390,4 +390,7 @@ def finetune(net, data_loader, val_data_loader, GPUS, epochs, eval_period, save_
         drop_last=True,
     )
 
+    print("Finish data loading...")
+
+    print("Start finetuning...")
     finetune(net, fine_dataloader, fine_val_dataloader, GPUS, 100, 20, 20)