107 lines
5.0 KiB
Python
107 lines
5.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Train and export the binary building-proposal filter."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
def binary_metrics(scores: list[float], labels: list[int], threshold: float = 0.5) -> dict[str, float | int]:
|
|
tp = sum(score >= threshold and label == 1 for score, label in zip(scores, labels, strict=True))
|
|
fp = sum(score >= threshold and label == 0 for score, label in zip(scores, labels, strict=True))
|
|
fn = sum(score < threshold and label == 1 for score, label in zip(scores, labels, strict=True))
|
|
precision = tp / (tp + fp) if tp + fp else 1.0
|
|
recall = tp / (tp + fn) if tp + fn else 1.0
|
|
return {"tp": tp, "fp": fp, "fn": fn, "precision": precision, "recall": recall,
|
|
"f1": 2 * precision * recall / (precision + recall) if precision + recall else 0.0}
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--dataset-dir", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--epochs", type=int, default=12)
|
|
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() and not args.export_existing_best:
|
|
parser.error(f"output already exists: {args.output_dir}")
|
|
|
|
import torch
|
|
from torch import nn
|
|
from torch.utils.data import DataLoader
|
|
from torchvision.datasets import ImageFolder
|
|
from torchvision.models import ResNet18_Weights, resnet18
|
|
|
|
weights = ResNet18_Weights.DEFAULT
|
|
transform = weights.transforms()
|
|
train_ds = ImageFolder(args.dataset_dir / "train", transform=transform)
|
|
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}")
|
|
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))
|
|
optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
|
|
args.output_dir.mkdir(parents=True)
|
|
history = []
|
|
best_f1 = -1.0
|
|
for epoch in range(1, args.epochs + 1):
|
|
model.train()
|
|
train_loss = 0.0
|
|
for images, labels in train_loader:
|
|
images, labels = images.to(device), labels.float().to(device)
|
|
optimizer.zero_grad(set_to_none=True)
|
|
logits = model(images).flatten()
|
|
loss = loss_fn(logits, labels)
|
|
loss.backward()
|
|
optimizer.step()
|
|
train_loss += float(loss) * len(images)
|
|
model.eval()
|
|
scores: list[float] = []
|
|
labels_out: list[int] = []
|
|
with torch.inference_mode():
|
|
for images, labels in val_loader:
|
|
scores.extend(torch.sigmoid(model(images.to(device)).flatten()).cpu().tolist())
|
|
labels_out.extend(labels.tolist())
|
|
metric = binary_metrics(scores, labels_out)
|
|
row = {"epoch": epoch, "train_loss": train_loss / len(train_ds), **metric}
|
|
history.append(row)
|
|
print(json.dumps(row), flush=True)
|
|
if float(metric["f1"]) > best_f1:
|
|
best_f1 = float(metric["f1"])
|
|
torch.save(model.state_dict(), args.output_dir / "best-state.pt")
|
|
model.load_state_dict(torch.load(args.output_dir / "best-state.pt", map_location=device))
|
|
model.eval()
|
|
scripted = torch.jit.script(model)
|
|
scripted.save(str(args.output_dir / "proposal-classifier.torchscript.pt"))
|
|
report = {"schema_version": 1, "status": "ok", "classes": train_ds.class_to_idx,
|
|
"train_count": len(train_ds), "validation_count": len(val_ds),
|
|
"best_validation_f1": best_f1, "history": history,
|
|
"model": str(args.output_dir / "proposal-classifier.torchscript.pt")}
|
|
(args.output_dir / "training-report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|