Frontend 5/5 speedup: Frontend CPU/GPU overlap

vidmap/frontend/keyframes/matching.py

@@ -1,5 +1,7 @@
 """Ordered low-resolution matching for keyframe selection."""
 
+from contextlib import contextmanager
+
 import torch
 from torch.utils.data import DataLoader
 
@@ -71,23 +73,51 @@ def build_pair_loader(scene_parser, sequence, lowres_options):
     return loader, len(pair_indices), original_width, original_height
 
 
-def match_lowres_batch(tracker_model, batch, original_width, original_height, first_batch, *, batch_size):
+@contextmanager
+def pipelined_matches(tracker_model, loader, original_width, original_height, *, batch_size):
+    """Overlap one ordered inference batch with CPU selection, draining on exit."""
     from vidmap.utils.profiling import record_timing, sync_time
 
-    names_A = batch["names_A"]
-    names_B = batch["names_B"]
-    batch_start = sync_time()
-    im_A = batch["im_A_batch"]
-    im_B = batch["im_B_batch"]
-    output = tracker_model.match_lowres_batch(
-        im_A,
-        im_B,
-        names_a=names_A,
-        names_b=names_B,
-        output_size=(original_width, original_height),
-        batch_size=batch_size,
-    )
-    if first_batch:
-        record_timing("keyframing_first_batch", sync_time() - batch_start, first=True)
-
-    return output.matches, output.certainty
+    source = iter(loader)
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    first_batch = True
+
+    def launch(batch):
+        nonlocal first_batch
+        with torch.cuda.stream(stream):
+            batch_start = sync_time()
+            output = tracker_model.match_lowres_batch(
+                batch["im_A_batch"],
+                batch["im_B_batch"],
+                names_a=batch["names_A"],
+                names_b=batch["names_B"],
+                output_size=(original_width, original_height),
+                batch_size=batch_size,
+            )
+            if first_batch:
+                record_timing("keyframing_first_batch", sync_time() - batch_start, first=True)
+            ready = torch.cuda.Event()
+            ready.record()
+        first_batch = False
+        return (output.matches, output.certainty), ready
+
+    def iterate():
+        first = next(source, None)
+        pending = None if first is None else launch(first)
+        while pending is not None:
+            (matches, certainties), ready = pending
+            consumer = torch.cuda.current_stream()
+            consumer.wait_event(ready)
+            matches.record_stream(consumer)
+            certainties.record_stream(consumer)
+            following = next(source, None)
+            pending = None if following is None else launch(following)
+            yield matches, certainties
+
+    batches = iterate()
+    try:
+        yield batches
+    finally:
+        stream.synchronize()
+        batches.close()

vidmap/frontend/keyframes/processing.py

@@ -180,22 +180,21 @@ def select_candidates(self, salient_metadata) -> KeyframeCandidates:
                 bootstrap_intrinsics,
             )
 
-            first_batch = True
-            with tqdm(
-                total=total_pairs,
-                desc="Streaming keyframe detection",
-                disable=not progress_bars_enabled(),
-            ) as progress:
-                for batch in loader:
-                    matches, certainties = keyframe_matching.match_lowres_batch(
-                        tracker_model,
-                        batch,
-                        original_width,
-                        original_height,
-                        first_batch,
-                        batch_size=self.lowres_options.batch_size,
-                    )
-                    first_batch = False
+            with (
+                keyframe_matching.pipelined_matches(
+                    tracker_model,
+                    loader,
+                    original_width,
+                    original_height,
+                    batch_size=self.lowres_options.batch_size,
+                ) as batches,
+                tqdm(
+                    total=total_pairs,
+                    desc="Streaming keyframe detection",
+                    disable=not progress_bars_enabled(),
+                ) as progress,
+            ):
+                for matches, certainties in batches:
                     for pair_match_lr, pair_cert_lr in zip(matches, certainties):
                         selector.process_pair(pair_match_lr, pair_cert_lr)
                         progress.update(1)

vidmap/frontend/options/matching.py

@@ -28,7 +28,7 @@ class PreprocessingOptions:
 @dataclass(frozen=True, config=ConfigDict(extra="forbid"))
 class LowresMatchOptions:
     batch_size: int = 8
-    num_workers: int = 16
+    num_workers: int = 4
     resize_to_shape: Tuple[int, int] = (560, 560)
     interpolation: str = "torch_bicubic"
 

vidmap/frontend/tracking/composition.py

@@ -1,6 +1,7 @@
 """Concrete composition of frontend's tracking-owned stages."""
 
 import logging
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 
 from vidmap.frontend.cache import fingerprint, ordered_files_fingerprint
@@ -46,6 +47,10 @@ def __init__(
 
     def run(self):
         """Run model stages, then depth and optional camera calibration."""
+        with ThreadPoolExecutor(max_workers=1, thread_name_prefix="da3-verify") as preparation:
+            return self._run(preparation)
+
+    def _run(self, preparation):
         from vidmap.frontend.pipeline import FrontendArtifacts, TrackingFrontendResult, validate_tracking_result
 
         options = self.options
@@ -71,6 +76,7 @@ def run(self):
             logger.info("Geo-calibration (batch): %s", paths.geocalib_batch_path)
         frames = FrameSequence.from_scene(self.scene_parser)
 
+        verification = None
         with create_lazy_romav2_tracker(tracker_options) as tracker:
             keyframe_processor = KeyframeProcessor(
                 scene_parser=self.scene_parser,
@@ -116,6 +122,9 @@ def run(self):
                     track_pairs_metadata=track_pairs_metadata,
                 )
                 candidates = keyframe_processor.select_candidates(provisional_salient_metadata)
+                from vidmap.frontend.models.depth.da3_video import verify_da3_model_snapshot
+
+                verification = preparation.submit(verify_da3_model_snapshot)
                 keyframes, tracks = sparse_track_builder.admit_and_build_tracks(
                     candidates,
                     options.keyframes.selection if options.keyframes.selection.lookahead_pruning else None,
@@ -151,6 +160,8 @@ def run(self):
             )
             extended = extended_match_builder.build()
 
+        if verification is not None:
+            verification.result()
         depth_estimator = DepthEstimator(
             scene_parser=self.scene_parser,
             paths=paths,