From 814f336bda0295c3fa100a4dd89b0b7746eb61f8 Mon Sep 17 00:00:00 2001
From: Hugues Sibille <hugues.sibille@epfl.ch>
Date: Tue, 25 Aug 2026 11:32:32 +0200
Subject: [PATCH] perf: keep the LoGeR input batch on the host

Pi3.forward already moves each sliding window to the model's device, so
preloading the whole resized sequence onto the GPU held VRAM for frames
that were not being decoded. Dropping the preload saves 1.6 MiB of VRAM
per frame at the default 504x280 processing size.

Measured on an RTX 3090 over 600 frames: peak allocated 12561 -> 11587
MiB, inference 59.5 -> 58.8 s, and a checksum over depth, poses, points
and confidence is unchanged.

The MPS fallback's batch_t.cpu() is now a no-op and is removed; moving
the model is what that path needs.
---
 CHANGELOG.md                         | 11 +++++++
 deepreefmap/mapping/loger_backend.py |  6 ++--
 tests/test_loger_adapter_contract.py | 46 ++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bac1027..2c57ec7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,17 @@ affect published measurements.
 
 ## [Unreleased]
 
+### Fixed
+
+- LoGeR peak VRAM drops by the size of the input batch, 1.6 MiB per frame at the
+  default 504x280 processing size. The backend preloaded the whole resized
+  sequence onto the GPU, but `Pi3.forward` already moves each sliding window to
+  the model's device itself, so only the window being decoded ever needed to be
+  resident. Measured on an RTX 3090 over 600 frames: 12561 MiB to 11587 MiB peak
+  allocated, a 974 MiB saving, with inference time unchanged (59.5 s to 58.8 s).
+  Outputs are byte-identical, checked with a checksum over depth, poses, points
+  and confidence. (#19)
+
 ## [1.1.0] - 2026-08-21
 
 ### Added
diff --git a/deepreefmap/mapping/loger_backend.py b/deepreefmap/mapping/loger_backend.py
index 73ac3f5..941c286 100644
--- a/deepreefmap/mapping/loger_backend.py
+++ b/deepreefmap/mapping/loger_backend.py
@@ -214,7 +214,10 @@ def process_sequence(
             )
             batch = np.stack(resized, axis=0).astype(np.float32) / 255.0
             del resized  # only needed to build the batch; frees ~3 B/px/frame during inference
-            batch_t = torch.from_numpy(batch).permute(0, 3, 1, 2).unsqueeze(0).to(self._device)
+            # Stays on the host: Pi3.forward moves each window to the model's device
+            # itself, so preloading the sequence only pinned VRAM for frames that are
+            # not being decoded yet.
+            batch_t = torch.from_numpy(batch).permute(0, 3, 1, 2).unsqueeze(0)
             del batch
             forward_kwargs = {
                 "window_size": self.default_window_size,
@@ -251,7 +254,6 @@ def process_sequence(
                             raise
                         logger.warning("LoGeR op unsupported on MPS, retrying on CPU: %s", exc)
                         model = model.cpu()
-                        batch_t = batch_t.cpu()
                         self._device = torch.device("cpu")
                         out = model(batch_t, **forward_kwargs)
             finally:
diff --git a/tests/test_loger_adapter_contract.py b/tests/test_loger_adapter_contract.py
index f702794..8b1836a 100644
--- a/tests/test_loger_adapter_contract.py
+++ b/tests/test_loger_adapter_contract.py
@@ -19,6 +19,52 @@ def test_loger_disables_per_frame_proxy_path():
         backend.process_frame(0, np.zeros((4, 4, 3), dtype=np.uint8))
 
 
+def test_loger_hands_the_model_a_host_resident_batch(monkeypatch):
+    """Pi3.forward moves each window to the model's device, so the sequence stays on the host."""
+    import torch
+
+    seen: dict[str, object] = {}
+    moved: list[object] = []
+    original_to = torch.Tensor.to
+
+    def recording_to(self, *args, **kwargs):
+        if self.dim() == 5:  # the (B, N, C, H, W) input batch, nothing else
+            moved.append(args[0] if args else kwargs.get("device"))
+        return original_to(self, *args, **kwargs)
+
+    def fake_model(imgs, **kwargs):
+        seen["device"] = imgs.device
+        seen["shape"] = tuple(imgs.shape)
+        n = imgs.shape[1]
+        return {
+            "local_points": torch.zeros(n, 4, 4, 3),
+            "points": torch.zeros(n, 4, 4, 3),
+            "camera_poses": torch.eye(4).repeat(n, 1, 1),
+            "conf": torch.zeros(n, 4, 4, 1),
+        }
+
+    backend = LoGeRBackend.__new__(LoGeRBackend)
+    backend._torch = torch
+    backend._model = fake_model
+    backend._device = torch.device("cpu")
+    backend._target_resolution = (56, 56)
+    backend.default_window_size = 2
+    backend._overlap_size = 0
+    backend._config = {}
+    backend._se3 = backend._sim3 = False
+    backend._turn_off_ttt = backend._turn_off_swa = False
+    backend._image_size = (8, 8)
+    backend._k = np.eye(3, dtype=np.float32)
+
+    frames = [np.zeros((8, 8, 3), dtype=np.uint8) for _ in range(3)]
+    monkeypatch.setattr(torch.Tensor, "to", recording_to)
+    backend.process_sequence([0, 1, 2], frames)
+
+    assert seen["device"] == torch.device("cpu")
+    assert seen["shape"] == (1, 3, 3, 56, 56)
+    assert moved == []
+
+
 def test_loger_target_resolution_uses_patch_multiple():
     assert _nearest_multiple(448, 14) == 448
     assert _nearest_multiple(450, 14) == 448