perf(inference): widen uint8 predict() inputs on the accelerator

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.

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

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()