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