From c9213822328feaff14e272869f97055eb9ed3848 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jes=C3=BAs=20Royeth?= Date: Mon, 31 Aug 2026 13:59:39 -0400 Subject: [PATCH] perf(inference): widen uint8 predict() inputs on the accelerator --- CHANGELOG.md | 2 + src/rfdetr/detr.py | 68 +++++++++++++++++---- tests/inference/test_predict.py | 101 ++++++++++++++++++++++++++++++-- 3 files changed, 153 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c2124c79e..79ba8f385 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ### Changed +- `RFDETR.predict()` now carries PIL and uint8 NumPy inputs across the host-to-device boundary in their original one-byte-per-channel storage and widens them to the default floating-point dtype on the model's device, instead of widening on the host and transferring four times the bytes. The dtype/layout fusion itself is unchanged — only where it runs — so a CPU model behaves exactly as before. The divisor is now a 0-dim tensor rather than the Python `int` `255`, because CUDA evaluates `Tensor.div_(255)` as a multiplication by the scalar's reciprocal and would otherwise land 1 ULP away from the host result for just under half of all byte values (126 of 256); with a tensor divisor the produced pixels stay byte-for-byte identical to `torchvision.transforms.functional.to_tensor`'s output on every device this change was verified on (CPU and CUDA). This is the host-to-device counterpart of the byte transfer the CUDA source-image path already performs. + - `RFDETR.predict()` now converts PIL and uint8 NumPy inputs from HWC byte storage to contiguous CHW floating-point storage with one dtype/layout allocation and in-place scaling. The default source-image path reuses its already-materialized PIL array, while non-uint8 NumPy inputs retain torchvision's conversion path. The `[0, 1]` range-scan skip for those same two input types is unaffected: the fused conversion divides `uint8` storage by 255 and carries the same guarantee `to_tensor` did. - The deformable-attention core now reuses its sampled tensor directly for one-level inputs instead of stacking and flattening a one-element list. Multi-level packing is unchanged. This primarily removes allocation work from keypoint cross-attention, whose current configuration uses one feature level and many more queries than the detection decoder. diff --git a/src/rfdetr/detr.py b/src/rfdetr/detr.py index 71de03648..1750544fc 100644 --- a/src/rfdetr/detr.py +++ b/src/rfdetr/detr.py @@ -92,30 +92,58 @@ def _tensor_to_source_array(image: torch.Tensor) -> np.ndarray[Any, Any]: return (source_view.cpu().numpy() * 255).astype(np.uint8) -def _uint8_image_to_tensor(image: np.ndarray[Any, Any]) -> torch.Tensor: - """Convert a 2-D/3-D uint8 image to contiguous CHW floating-point storage. +def _uint8_image_to_chw_view(image: np.ndarray[Any, Any]) -> torch.Tensor: + """Return a zero-copy ``uint8`` CHW *view* of a 2-D/3-D HWC image array. - This is the uint8 branch of ``torchvision.transforms.functional.to_tensor`` with its layout/dtype conversion fused - and division performed on fresh storage in place. + This is the layout half of ``torchvision.transforms.functional.to_tensor``; the dtype half is + :func:`_uint8_chw_to_float`. Splitting them lets :meth:`RFDETR.predict` send the + 1-byte-per-channel storage across the host-to-device boundary and widen it on the accelerator, + instead of widening it 4x on the host and transferring that. Args: image: A ``(H, W)`` grayscale or ``(H, W, C)`` HWC ``uint8`` array. Returns: - A contiguous ``(C, H, W)`` tensor in the current default floating-point dtype, scaled to ``[0, 1]``, - byte-for-byte identical to ``to_tensor``'s output. For ``C == 1`` the leading dimension's own stride - may differ from ``to_tensor``'s, which is immaterial: a size-1 dimension's stride has no memory-layout - effect. + A ``(C, H, W)`` ``uint8`` tensor sharing *image*'s storage. Examples: >>> arr = np.zeros((2, 2, 3), dtype=np.uint8) - >>> _uint8_image_to_tensor(arr).shape + >>> _uint8_image_to_chw_view(arr).shape torch.Size([3, 2, 2]) """ if image.ndim == 2: image = image[:, :, None] - chw = torch.from_numpy(image.transpose((2, 0, 1))) - return chw.to(dtype=torch.get_default_dtype(), memory_format=torch.contiguous_format).div_(255) + return torch.from_numpy(image.transpose((2, 0, 1))) + + +def _uint8_chw_to_float(chw: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Widen a ``uint8`` CHW tensor to the default float dtype and scale it into ``[0, 1]``. + + Dtype and layout are converted together and the division is done in place on that fresh float + allocation, so the pair costs one allocation rather than ``to_tensor``'s three. + + Args: + chw: ``(C, H, W)`` ``uint8`` tensor, on any device. + scale: 0-dim tensor holding ``255``, on ``chw``'s device and in the target dtype. It has to + be a tensor rather than the Python ``int``: CUDA evaluates ``Tensor.div_(255)`` as a + multiplication by ``255``'s reciprocal, which rounds differently from the host's + division for just under half of all byte values (126 of 256, 1 ULP). A tensor divisor + keeps IEEE-754 correctly rounded division on both, so the result does not depend on + where it was computed. + + Returns: + A contiguous ``(C, H, W)`` tensor in the current default floating-point dtype, scaled to + ``[0, 1]``, byte-for-byte identical to ``to_tensor``'s output. For ``C == 1`` the leading + dimension's own stride may differ from ``to_tensor``'s, which is immaterial: a size-1 + dimension's stride has no memory-layout effect. + + Examples: + >>> chw = _uint8_image_to_chw_view(np.zeros((2, 2, 3), dtype=np.uint8)) + >>> _uint8_chw_to_float(chw, torch.tensor(255.0)).dtype + torch.float32 + """ + widened = chw.to(dtype=torch.get_default_dtype(), memory_format=torch.contiguous_format) + return widened.div_(scale) # ModelContext and _build_model_context are eagerly imported above (runtime use in get_model). @@ -2440,6 +2468,8 @@ class and ``class_id=1`` is ``"__background__"``. # original code checked range before shape for a given image, and raising it eagerly here # would flip that precedence for any tensor that is invalid on both axes at once. pending_checks: list[tuple[torch.Tensor | bool, torch.Tensor | bool, bool, tuple[int, ...]]] = [] + # Built lazily on the first uint8 image, then shared by the rest of the batch. + uint8_scale: torch.Tensor | None = None for img_input in images: img: Any = img_input @@ -2451,6 +2481,7 @@ class and ``class_id=1`` is ``"__background__"``. img = Image.open(img) range_known_valid = False + deferred_widen = False if not isinstance(img, torch.Tensor): # Auto-convert PIL images from any colour mode (L, LA, RGBA, P, # etc.) to RGB before converting to tensor. This matches the @@ -2480,7 +2511,13 @@ class and ``class_id=1`` is ``"__background__"``. ) else: tensor_source = img - img = _uint8_image_to_tensor(tensor_source) + # Keep the 1-byte-per-channel storage for now: the widening to float is + # deferred until after the host-to-device transfer below, so only a quarter of + # the bytes cross the bus and the widen+divide run on the accelerator. The view + # is already (C, H, W), so every shape check and error message below is + # unchanged. + img = _uint8_image_to_chw_view(tensor_source) + deferred_widen = True else: img = F.to_tensor(img) elif include_source_image and img.dim() == 3: @@ -2530,7 +2567,12 @@ class and ``class_id=1`` is ``"__background__"``. # CPU-model transfer with non_blocking=True races the copy — the CPU destination is never pinned, so # reads of the tensor's data can observe an in-flight (partially written) copy. non_blocking = self.model.device.type == "cuda" - processed_images.append(img_tensor.to(self.model.device, non_blocking=non_blocking)) + img_tensor = img_tensor.to(self.model.device, non_blocking=non_blocking) + if deferred_widen: + if uint8_scale is None: + uint8_scale = torch.tensor(255, device=img_tensor.device, dtype=torch.get_default_dtype()) + img_tensor = _uint8_chw_to_float(img_tensor, uint8_scale) + processed_images.append(img_tensor) # Force the range-check results to Python bools only now, after every image's conversion, # range-check kernels, and transfer have all been queued (see the comment where diff --git a/tests/inference/test_predict.py b/tests/inference/test_predict.py index e9d152709..33d38b142 100644 --- a/tests/inference/test_predict.py +++ b/tests/inference/test_predict.py @@ -753,7 +753,9 @@ def test_converter_is_bit_exact_for_supported_default_dtypes(self, dtype: torch. try: torch.set_default_dtype(dtype) expected = F.to_tensor(image) - actual = detr_module._uint8_image_to_tensor(image) + actual = detr_module._uint8_chw_to_float( + detr_module._uint8_image_to_chw_view(image), torch.tensor(255, dtype=dtype) + ) finally: torch.set_default_dtype(previous_dtype) @@ -796,7 +798,7 @@ def test_predict_uses_fused_path_for_single_channel_uint8_images(self, image: np to_tensor_spy = MagicMock(side_effect=AssertionError("uint8 input unexpectedly used F.to_tensor")) with ( - patch("rfdetr.detr._uint8_image_to_tensor", wraps=detr_module._uint8_image_to_tensor) as converter_spy, + patch("rfdetr.detr._uint8_image_to_chw_view", wraps=detr_module._uint8_image_to_chw_view) as converter_spy, patch("rfdetr.detr.F.to_tensor", to_tensor_spy), ): model.predict(image) @@ -816,7 +818,7 @@ def test_converter_matches_torchvision_for_contiguous_single_channel_hwc(self) - image = rng.integers(0, 256, size=(17, 29, 1), dtype=np.uint8) expected = F.to_tensor(image) - actual = detr_module._uint8_image_to_tensor(image) + actual = detr_module._uint8_chw_to_float(detr_module._uint8_image_to_chw_view(image), torch.tensor(255.0)) assert actual.dtype == expected.dtype assert actual.shape == expected.shape @@ -829,7 +831,7 @@ def test_predict_reuses_pil_source_array_for_conversion(self) -> None: model = _DummyRFDETR() image = PIL.Image.new("RGB", (29, 17), color=(1, 127, 255)) - with patch("rfdetr.detr._uint8_image_to_tensor", wraps=detr_module._uint8_image_to_tensor) as converter_spy: + with patch("rfdetr.detr._uint8_image_to_chw_view", wraps=detr_module._uint8_image_to_chw_view) as converter_spy: detections = model.predict(image) converter_spy.assert_called_once() @@ -843,7 +845,7 @@ def test_predict_keeps_original_readonly_numpy_as_tensor_source(self) -> None: image.flags.writeable = False with ( - patch("rfdetr.detr._uint8_image_to_tensor", wraps=detr_module._uint8_image_to_tensor) as converter_spy, + patch("rfdetr.detr._uint8_image_to_chw_view", wraps=detr_module._uint8_image_to_chw_view) as converter_spy, pytest.warns(UserWarning, match="not writable"), ): detections = model.predict(image) @@ -853,6 +855,95 @@ def test_predict_keeps_original_readonly_numpy_as_tensor_source(self) -> None: assert detections.metadata["source_image"] is not image assert detections.metadata["source_image"].flags.writeable + def test_deferred_widening_is_bit_exact_for_every_byte_value(self) -> None: + """The 0-dim divisor keeps IEEE-754 division, which a Python scalar does not on CUDA. + + ``Tensor.div_(255)`` is evaluated on CUDA as a multiplication by ``1/255``; that rounds differently from the + CPU's division for just under half of all byte values (126 of 256). Since ``predict()`` now widens after the + transfer, the divisor must be a 0-dim tensor for the accelerator result to stay byte-for-byte equal to + torchvision's host-side conversion. + """ + image = np.arange(256, dtype=np.uint8).reshape(16, 16, 1).repeat(3, axis=2) + chw = detr_module._uint8_image_to_chw_view(image) + + assert detr_module._uint8_chw_to_float(chw, torch.tensor(255.0)).equal(F.to_tensor(image)) + + @pytest.mark.gpu + def test_deferred_widening_is_bit_exact_on_real_cuda(self) -> None: + """The reciprocal-multiplication gap this fix avoids only manifests on real CUDA hardware. + + ``test_deferred_widening_is_bit_exact_for_every_byte_value`` above proves the divisor swap is correct on + CPU, where scalar and tensor division already agree bit-for-bit; it cannot exercise the CUDA-specific + ``Tensor.div_(255)`` reciprocal path this fix replaces. This test runs the same view/pin/transfer/widen + sequence ``predict()`` uses on a real CUDA device and compares it against the host-computed reference. + """ + image = np.arange(256, dtype=np.uint8).reshape(16, 16, 1).repeat(3, axis=2) + expected = F.to_tensor(image) + + chw = detr_module._uint8_image_to_chw_view(image).pin_memory().to("cuda", non_blocking=True) + scale = torch.tensor(255, device=chw.device, dtype=torch.get_default_dtype()) + actual = detr_module._uint8_chw_to_float(chw, scale).cpu() + + assert actual.equal(expected) + + def test_one_divisor_serves_every_image_of_a_multi_image_call(self) -> None: + """`predict()` builds the divisor once per call and reuses it for the rest of the batch. + + The in-place division has to consume the freshly widened copy, not the shared divisor, or the second and later + images of a multi-image call would be scaled by an already-divided value. + """ + scale = torch.tensor(255, dtype=torch.get_default_dtype()) + dark = np.full((4, 5, 3), 51, dtype=np.uint8) + bright = np.full((4, 5, 3), 255, dtype=np.uint8) + + first = detr_module._uint8_chw_to_float(detr_module._uint8_image_to_chw_view(dark), scale) + second = detr_module._uint8_chw_to_float(detr_module._uint8_image_to_chw_view(bright), scale) + + assert first.equal(F.to_tensor(dark)) + assert second.equal(F.to_tensor(bright)) + assert scale.equal(torch.tensor(255, dtype=torch.get_default_dtype())) + + def test_uint8_images_cross_the_device_boundary_unwidened(self) -> None: + """Only the 1-byte-per-channel storage is transferred; the widening runs after the copy.""" + model = _DummyRFDETR() + model.model.device = torch.device("cuda", 0) + real_to = torch.Tensor.to + real_tensor = torch.tensor + transferred_dtypes: list[torch.dtype] = [] + + def to_spy(self: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor: + """Record the dtype of each simulated CUDA transfer and keep the result on CPU. + + Examples: + This test-local closure requires its enclosing capture list. + >>> to_spy(torch.zeros(3), torch.device("cuda")) # doctest: +SKIP + """ + if args and isinstance(args[0], torch.device) and args[0].type == "cuda": + transferred_dtypes.append(self.dtype) + return self + return real_to(self, *args, **kwargs) + + def tensor_spy(data: object, **kwargs: object) -> torch.Tensor: + """Redirect CUDA target-size construction to CPU for the simulated transfer. + + Examples: + This test-local closure requires the captured real tensor constructor. + >>> tensor_spy([[48, 64]], device=torch.device("cuda")) # doctest: +SKIP + """ + device = kwargs.get("device") + if isinstance(device, torch.device) and device.type == "cuda": + kwargs["device"] = torch.device("cpu") + return real_tensor(data, **kwargs) + + with ( + patch.object(torch.Tensor, "pin_memory", lambda self: self), + patch.object(torch.Tensor, "to", to_spy), + patch.object(torch, "tensor", tensor_spy), + ): + model.predict(np.full((17, 29, 3), 127, dtype=np.uint8)) + + assert transferred_dtypes == [torch.uint8] + def test_predict_keeps_torchvision_fallback_for_float_numpy(self) -> None: """Non-uint8 NumPy inputs retain torchvision's no-scaling conversion semantics.""" model = _DummyRFDETR()