fix dry run

run_cspca.py

@@ -21,8 +21,9 @@ def main_worker(args):
     cache_dir_path = Path(os.path.join(args.logdir, "cache"))
 
     if args.mode == "train":
-        checkpoint = torch.load(args.checkpoint_pirads, weights_only=False, map_location="cpu")
-        mil_model.load_state_dict(checkpoint["state_dict"])
+        if not args.dry_run:
+            checkpoint = torch.load(args.checkpoint_pirads, weights_only=False, map_location="cpu")
+            mil_model.load_state_dict(checkpoint["state_dict"])
         mil_model = mil_model.to(args.device)
 
         model_dir = os.path.join(args.logdir, "models")
@@ -65,8 +66,11 @@ def main_worker(args):
 
 
     cspca_model = CSPCAModel(backbone=mil_model).to(args.device)
-    checkpt = torch.load(args.checkpoint_cspca, map_location="cpu")
-    cspca_model.load_state_dict(checkpt["state_dict"])
+
+    if not args.dry_run:
+        checkpt = torch.load(args.checkpoint_cspca, map_location="cpu")
+        cspca_model.load_state_dict(checkpt["state_dict"])
+
     cspca_model = cspca_model.to(args.device)
     if "auc" in checkpt and "sensitivity" in checkpt and "specificity" in checkpt:
         auc, sens, spec = checkpt["auc"], checkpt["sensitivity"], checkpt["specificity"]
@@ -191,12 +195,13 @@ if __name__ == "__main__":
     if args.dataset_json is None:
         logging.error("Dataset path not provided. Quitting.")
         sys.exit(1)
-    if args.checkpoint_pirads is None and args.mode == "train":
-        logging.error("PI-RADS checkpoint path not provided. Quitting.")
-        sys.exit(1)
-    elif args.checkpoint_cspca is None and args.mode == "test":
-        logging.error("csPCa checkpoint path not provided. Quitting.")
-        sys.exit(1)
+    if not args.dry_run:
+        if args.checkpoint_pirads is None and args.mode == "train":
+            logging.error("PI-RADS checkpoint path not provided. Quitting.")
+            sys.exit(1)
+        elif args.checkpoint_cspca is None and args.mode == "test":
+            logging.error("csPCa checkpoint path not provided. Quitting.")
+            sys.exit(1)
     args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     if args.device == torch.device("cuda"):

run_pirads.py

@@ -26,7 +26,7 @@ def main_worker(args):
     model = MILModel3D(num_classes=args.num_classes, mil_mode=args.mil_mode)
     start_epoch = 0
     best_acc = 0.0
-    if args.checkpoint is not None:
+    if args.checkpoint is not None and not args.dry_run:
         checkpoint = torch.load(args.checkpoint, map_location="cpu")
         model.load_state_dict(checkpoint["state_dict"])
 
@@ -285,9 +285,10 @@ if __name__ == "__main__":
     if args.dataset_json is None:
         logging.error("Dataset JSON file not provided. Quitting.")
         sys.exit(1)
-    if args.checkpoint is None and args.mode == "test":
-        logging.error("Model checkpoint path not provided. Quitting.")
-        sys.exit(1)
+    if not args.dry_run:
+        if args.checkpoint is None and args.mode == "test":
+            logging.error("Model checkpoint path not provided. Quitting.")
+            sys.exit(1)
 
     if args.dry_run:
         logging.info("Dry run mode enabled.")