diff --git a/CHANGELOG.md b/CHANGELOG.md index c2124c79e..d7eb93746 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -62,6 +62,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), ### Fixed +- Packed targets are now materialized directly into independent per-sample tensors on the destination device instead of first moving every concatenated field and then cloning it. This removes a transient CUDA allocation equal to the mask field's size, present whenever segmentation transfer held both the packed and the materialized masks at once; DataLoader worker transport remains packed. This is a mechanism-level fix, not the representative detection/segmentation host-device memory benchmark suite noted as a follow-up boundary above. + - Empty COCO targets now keep `iscrowd` as `int64` and `area` as `float32`, matching populated targets and allowing mixed empty/populated batches to use lossless packed-target worker transport. --- diff --git a/docs/learn/train/training-parameters.md b/docs/learn/train/training-parameters.md index 38239943e..c8e10a919 100644 --- a/docs/learn/train/training-parameters.md +++ b/docs/learn/train/training-parameters.md @@ -301,7 +301,7 @@ The parameters below are available for fine-grained control over training behavi | `prefetch_factor` | `int` | `None` | Number of batches to prefetch per DataLoader worker. `None` uses PyTorch's built-in default. | | `pack_targets` | `bool` | `True` | Concatenate target dicts before crossing the DataLoader worker boundary. See the contract below; set `False` to opt out. | -With `pack_targets=True`, train, validation, test, and predict loaders yield batches whose target element is `PackedTargets` whenever packing is lossless. The Lightning `transfer_batch_to_device` hook accepts those batches or an unpacked tuple of target dicts. It moves packed fields to the target device, then materializes them into the same plain per-sample dict list that training, validation, test, and prediction hooks receive on the unpacked path. Batches that cannot be packed losslessly retain their original tuple of dicts. +With `pack_targets=True`, train, validation, test, and predict loaders yield batches whose target element is `PackedTargets` whenever packing is lossless. The Lightning `transfer_batch_to_device` hook accepts those batches or an unpacked tuple of target dicts. It materializes each packed field directly into its own independently owned per-sample tensor on the target device, producing the same plain per-sample dict list that training, validation, test, and prediction hooks receive on the unpacked path. Batches that cannot be packed losslessly retain their original tuple of dicts. ## Complete Parameter Reference diff --git a/src/rfdetr/training/module_data.py b/src/rfdetr/training/module_data.py index 38ad0b6fd..f9a084fff 100644 --- a/src/rfdetr/training/module_data.py +++ b/src/rfdetr/training/module_data.py @@ -814,9 +814,9 @@ def transfer_batch_to_device( non_blocking = device.type == "cuda" samples = samples.to(device, non_blocking=non_blocking) if isinstance(targets, PackedTargets): - # One transfer per field instead of one per field per sample, then materialised so that callers can - # mutate the dicts exactly as they do on the unpacked path. - targets = targets.to(device, non_blocking=non_blocking).as_list() + # Materialise directly on the destination so the packed and per-sample copies of a large field such as + # segmentation masks never coexist there. Each returned tensor owns its storage, matching the unpacked path. + targets = targets.to_list(device, non_blocking=non_blocking) else: targets = [{k: v.to(device, non_blocking=non_blocking) for k, v in t.items()} for t in targets] return samples, targets diff --git a/src/rfdetr/utilities/tensors.py b/src/rfdetr/utilities/tensors.py index 2c4a04e52..b6f78d21a 100644 --- a/src/rfdetr/utilities/tensors.py +++ b/src/rfdetr/utilities/tensors.py @@ -389,8 +389,9 @@ class PackedTargets: Sequence subclass would be taken apart and pinned tensor by tensor, rebuilding in the main process exactly what the packing avoided. Indexing rebuilds a dict of views: reassigning a key on the returned dict does not write back into the packed storage, but an in-place write into one of its tensors does, since the view shares - the packed field's storage. Call :meth:`as_list` for a materialised copy that is safe to mutate in place, the - way the unpacked batch already is. + the packed field's storage. Call :meth:`as_list` for a same-device materialised copy, or :meth:`to_list` to + materialise directly on another device; both return independent tensors that are safe to mutate in place, the way + the unpacked batch already is. Args: values: One flat tensor per field, holding that field concatenated across the batch. @@ -450,8 +451,8 @@ def pin_memory(self) -> PackedTargets: """Pin every packed field, keeping the batch packed. Called by PyTorch's pin-memory worker. Pinning the concatenated fields costs one pinned allocation - per field rather than one per field per sample, and leaves the batch in the packed form that - ``transfer_batch_to_device`` moves in a single copy per field. + per field rather than one per field per sample, and leaves the batch packed until + ``transfer_batch_to_device`` materialises its independent destination tensors. Returns: A packed batch whose fields are in pinned memory. @@ -483,6 +484,35 @@ def as_list(self) -> list[dict[str, Tensor]]: """ return [{key: value.clone() for key, value in self[position].items()} for position in range(self._batch)] + def to_list(self, device: torch.device | str, non_blocking: bool = False) -> list[dict[str, Tensor]]: + """Materialise independent per-sample tensors directly on *device*. + + Moving the packed object and then calling :meth:`as_list` makes the concatenated destination fields coexist + with all of their per-sample clones. That transient duplicate is negligible for boxes and labels but can be + large for segmentation masks. Transferring each packed view directly into its final independently-owned tensor + avoids the full-field destination copy while preserving the unpacked path's mutation semantics. + + Args: + device: Destination device. + non_blocking: Whether to request asynchronous copies. + + Returns: + One dict per sample, ordered as packed, with independent tensors on *device*. + """ + materialised: list[dict[str, Tensor]] = [{} for _ in range(self._batch)] + for key, flat in self._values.items(): + offset = 0 + for target, shape in zip(materialised, self._shapes[key], strict=True): + elements = _numel(shape) + value = flat[offset : offset + elements].reshape(shape) + target[key] = value.to( + device=device, + non_blocking=non_blocking, + copy=True, + ) + offset += elements + return materialised + def to(self, device: torch.device | str, non_blocking: bool = False) -> PackedTargets: """Move every packed field to *device* in one transfer per field. diff --git a/tests/training/test_module_data.py b/tests/training/test_module_data.py index bfb5da62e..dddaf8989 100644 --- a/tests/training/test_module_data.py +++ b/tests/training/test_module_data.py @@ -1325,6 +1325,24 @@ def test_packed_targets_are_unpacked_to_a_plain_list_on_target_device(self, fixt for key, value in original.items(): assert torch.equal(rebuilt[key], value) + def test_packed_targets_are_materialised_without_a_whole_batch_device_copy(self, fixture_training_setup) -> None: + """Packed transfer must not create a device copy of every field before constructing per-sample tensors.""" + _, _, dm = fixture_training_setup + samples, plain_targets = _make_batch() + packed_targets = pack_targets(plain_targets) + assert isinstance(packed_targets, PackedTargets) + + with patch.object( + PackedTargets, + "to", + side_effect=AssertionError("whole packed batch copied to the target device"), + ): + _, result_targets = dm.transfer_batch_to_device( + (samples, packed_targets), torch.device("cpu"), dataloader_idx=0 + ) + + assert isinstance(result_targets, list) + # --------------------------------------------------------------------------- # TestBackendResolution — validates augmentation_backend logic in setup("fit") diff --git a/tests/utilities/test_tensors.py b/tests/utilities/test_tensors.py index 48e13ff53..1085473e6 100644 --- a/tests/utilities/test_tensors.py +++ b/tests/utilities/test_tensors.py @@ -868,6 +868,90 @@ def test_as_list_tensors_do_not_alias_the_packed_storage(self) -> None: assert packed.fields["boxes"][0].item() == 1.0 assert torch.equal(packed.as_list()[0]["labels"], torch.tensor([3, 7])) + def test_to_list_same_device_does_not_alias_packed_storage(self) -> None: + """Direct materialisation must preserve the unpacked path's independent tensor ownership on CPU.""" + packed = pack_targets(self._batch()) + assert isinstance(packed, PackedTargets) + + materialised = packed.to_list(torch.device("cpu")) + materialised[0]["labels"][0] = 999 + materialised[0]["boxes"][0, 0] = 999.0 + + assert torch.equal(packed.fields["labels"], torch.tensor([3, 7, 5])) + assert packed.fields["boxes"][0].item() == 1.0 + + @pytest.mark.gpu + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_to_list_transfers_pinned_views_directly_to_cuda(self) -> None: + """Direct materialisation must preserve values and ownership across a non-blocking CUDA transfer.""" + packed = pack_targets(self._batch()) + assert isinstance(packed, PackedTargets) + pinned = packed.pin_memory() + + materialised = pinned.to_list(torch.device("cuda"), non_blocking=True) + torch.cuda.synchronize() + + assert materialised[0]["labels"].device.type == "cuda" + assert torch.equal(materialised[0]["labels"].cpu(), torch.tensor([3, 7])) + materialised[0]["labels"][0] = 999 + assert torch.equal(pinned.fields["labels"], torch.tensor([3, 7, 5])) + + @pytest.mark.gpu + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_to_list_bounds_the_transient_cuda_peak_for_a_mask_field(self) -> None: + """``to_list``'s docstring claims the avoided duplicate "can be large for segmentation masks" -- pin that claim + on the actual CUDA peak with a mask field, not just on values or ownership.""" + pinned = pack_targets( + [ + {"labels": torch.tensor([3, 7]), "masks": torch.ones((2, 256, 256), dtype=torch.bool)}, + {"labels": torch.tensor([5]), "masks": torch.ones((1, 256, 256), dtype=torch.bool)}, + ] + ).pin_memory() + mask_bytes = pinned.fields["masks"].numel() + + def peak_extra(materialise): + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + out = materialise() + torch.cuda.synchronize() + extra = torch.cuda.max_memory_allocated() - torch.cuda.memory_allocated() + del out + return extra + + old_extra = peak_extra(lambda: pinned.to(torch.device("cuda"), non_blocking=True).as_list()) + new_extra = peak_extra(lambda: pinned.to_list(torch.device("cuda"), non_blocking=True)) + + assert old_extra >= mask_bytes + assert new_extra < mask_bytes // 2 + + def test_to_returns_self_when_already_on_the_target_device(self) -> None: + """``to()`` keeps the batch packed instead of materialising it, unlike ``to_list()``. + + Losing its only caller in ``transfer_batch_to_device`` (replaced by ``to_list()``) must not leave it untested: a + no-op device request has to return the same instance, matching every no-op ``Tensor.to()`` call underneath it. + """ + packed = pack_targets(self._batch()) + assert isinstance(packed, PackedTargets) + + same_device = packed.to(torch.device("cpu")) + + assert same_device is packed + assert torch.equal(same_device.fields["labels"], torch.tensor([3, 7, 5])) + + @pytest.mark.gpu + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_to_moves_every_field_and_keeps_the_batch_packed_on_cuda(self) -> None: + """A genuine device change must move every field and return a new packed batch, not a materialised list.""" + packed = pack_targets(self._batch()) + assert isinstance(packed, PackedTargets) + + moved = packed.to(torch.device("cuda")) + + assert isinstance(moved, PackedTargets) + assert moved is not packed + assert all(tensor.device.type == "cuda" for tensor in moved.fields.values()) + assert torch.equal(moved.fields["labels"].cpu(), torch.tensor([3, 7, 5])) + def test_collate_packs_only_when_asked(self) -> None: """The collate contract only changes for callers that opt in.""" images = [torch.zeros(3, 8, 8), torch.zeros(3, 8, 8)]