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))