From 41e096799664de107a1937792314b1c27eb24094 Mon Sep 17 00:00:00 2001 From: Jens Date: Thu, 30 Jul 2026 01:47:15 +0200 Subject: [PATCH] Allow exporting an early-stopped proposal classifier --- scripts/train_building_proposal_classifier.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/scripts/train_building_proposal_classifier.py b/scripts/train_building_proposal_classifier.py index e9ff9986..542fecba 100644 --- a/scripts/train_building_proposal_classifier.py +++ b/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))