Allow exporting an early-stopped proposal classifier

scripts/train_building_proposal_classifier.py

@@ -26,8 +26,9 @@ def main() -> int:
     parser.add_argument("--batch", type=int, default=64)
     parser.add_argument("--lr", type=float, default=1e-4)
     parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--export-existing-best", action="store_true")
     args = parser.parse_args()
-    if args.output_dir.exists():
+    if args.output_dir.exists() and not args.export_existing_best:
         parser.error(f"output already exists: {args.output_dir}")
 
     import torch
@@ -42,12 +43,21 @@ def main() -> int:
     val_ds = ImageFolder(args.dataset_dir / "val", transform=transform)
     if train_ds.class_to_idx != {"negative": 0, "positive": 1}:
         raise RuntimeError(f"unexpected class order: {train_ds.class_to_idx}")
-    train_loader = DataLoader(train_ds, batch_size=args.batch, shuffle=True, num_workers=0)
-    val_loader = DataLoader(val_ds, batch_size=args.batch, shuffle=False, num_workers=0)
     device = torch.device(args.device)
     model = resnet18(weights=weights)
     model.fc = nn.Linear(model.fc.in_features, 1)
     model.to(device)
+    if args.export_existing_best:
+        state_path = args.output_dir / "best-state.pt"
+        if not state_path.is_file():
+            raise RuntimeError(f"missing existing best state: {state_path}")
+        model.load_state_dict(torch.load(state_path, map_location=device))
+        model.eval()
+        torch.jit.script(model).save(str(args.output_dir / "proposal-classifier.torchscript.pt"))
+        print(json.dumps({"status": "exported_existing_best", "model": str(state_path)}))
+        return 0
+    train_loader = DataLoader(train_ds, batch_size=args.batch, shuffle=True, num_workers=0)
+    val_loader = DataLoader(val_ds, batch_size=args.batch, shuffle=False, num_workers=0)
     positives = sum(label == 1 for _path, label in train_ds.samples)
     negatives = len(train_ds) - positives
     loss_fn = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([negatives / positives], device=device))