Bound proposal-filter inference memory

scripts/train_building_proposal_filter.py

@@ -84,8 +84,8 @@ def materialize_split(
         for label, boxes in (("positive", positives), ("negative", negatives)):
             directory = output_dir / split / label
             directory.mkdir(parents=True, exist_ok=True)
-            for box in boxes:
+            for box_index, box in enumerate(boxes):
-                path = directory / f"{tile['sample_slug']}_{tile['tile_index']:04d}_{counts[label]:07d}.png"
+                path = directory / f"{Path(tile['image_path']).stem}_{label}_{box_index:04d}.png"
                 expanded_crop(image, box, crop_scale).save(path)
                 counts[label] += 1
     return counts
@@ -121,6 +121,7 @@ def main() -> int:
     parser.add_argument("--negative-match-iou", type=float, default=0.05)
     parser.add_argument("--epochs", type=int, default=8)
     parser.add_argument("--batch", type=int, default=64)
+    parser.add_argument("--proposal-chunk-size", type=int, default=16)
     parser.add_argument("--seed", type=int, default=20260727)
     parser.add_argument("--force", action="store_true")
     args = parser.parse_args()
@@ -143,18 +144,27 @@ def main() -> int:
         for split in ("train", "val")
     }
     proposal_model = YOLO(str(args.proposal_model))
+    if args.proposal_chunk_size < 1:
+        raise SystemExit("--proposal-chunk-size must be positive")
     counts = {}
     for split, tiles in split_tiles.items():
-        results = proposal_model.predict(
+        counts[split] = {"positive": 0, "negative": 0}
-            [tile["image_path"] for tile in tiles], conf=args.proposal_confidence, imgsz=640,
+        for start in range(0, len(tiles), args.proposal_chunk_size):
-            max_det=1000, device=args.device, verbose=False,
+            chunk = tiles[start : start + args.proposal_chunk_size]
-        )
+            results = proposal_model.predict(
-        counts[split] = materialize_split(
+                [tile["image_path"] for tile in chunk], conf=args.proposal_confidence, imgsz=640,
-            tiles=tiles, results=results, output_dir=args.output_dir / "crops", split=split,
+                max_det=1000, device=args.device, verbose=False,
-            crop_scale=args.crop_scale, max_positives_per_tile=args.max_positives_per_tile,
+            )
-            max_negatives_per_tile=args.max_negatives_per_tile, negative_match_iou=args.negative_match_iou,
+            chunk_counts = materialize_split(
-        )
+                tiles=chunk, results=results, output_dir=args.output_dir / "crops", split=split,
-        del results
+                crop_scale=args.crop_scale, max_positives_per_tile=args.max_positives_per_tile,
+                max_negatives_per_tile=args.max_negatives_per_tile, negative_match_iou=args.negative_match_iou,
+            )
+            for label, count in chunk_counts.items():
+                counts[split][label] += count
+            del results
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
 
     del proposal_model
     if torch.cuda.is_available():
@@ -206,6 +216,7 @@ def main() -> int:
         "summary": str(args.summary), "summary_sha256": sha256(args.summary),
         "proposal_model": str(args.proposal_model), "proposal_model_sha256": sha256(args.proposal_model),
         "device": args.device, "proposal_confidence": args.proposal_confidence, "crop_scale": args.crop_scale,
+        "proposal_chunk_size": args.proposal_chunk_size,
         "negative_match_iou": args.negative_match_iou, "counts": counts, "epochs": args.epochs,
         "history": history, "selected_threshold_source": "validation_only",
         "selected_threshold": threshold, "validation_metrics": validation_metrics,