feat(streampetr): harden the streaming training infrastructure

.gitignore

@@ -2,13 +2,16 @@
 data/
 mlruns/
 outputs/
+work_dirs/
 
 # Model files and artifacts
 *.bin
 *.db
 *.onnx
 *.engine
 *.pth
+*.pt
+*.npz
 *.ckpt
 *.log
 *.out

autoware_ml/datamodule/base.py

@@ -19,7 +19,7 @@
 """
 
 from abc import ABC, abstractmethod
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
 from dataclasses import dataclass
 import logging
 from typing import Any
@@ -144,8 +144,14 @@ class DataModule(L.LightningDataModule, ABC):
 
     Collation contract:
         ``collation_map`` is an explicit whitelist of keys that may appear in
-        the collated batch. Every key listed must be present in every sample.
-        Missing keys are skipped with a warning. Each key declares exactly one strategy.
+        the collated batch of every split. ``train_collation_map`` declares
+        additional keys that only training samples carry (e.g. train-only
+        auxiliary supervision produced by train transforms); it is merged on
+        top of ``collation_map`` for the train dataloader only, and a key
+        present in both maps takes the train map's strategy. Keys listed for
+        a split are expected in every sample of that split; a declared key
+        missing from the samples is skipped with a warning. Each key declares
+        exactly one strategy.
 
         Strategies:
             * ``"concat"``: concatenate per-sample tensors along dim 0. The
@@ -169,6 +175,7 @@ class DataModule(L.LightningDataModule, ABC):
     def __init__(
         self,
         collation_map: Mapping[str, CollationStrategy] | None = None,
+        train_collation_map: Mapping[str, CollationStrategy] | None = None,
         train_transforms: TransformsCompose | None = None,
         val_transforms: TransformsCompose | None = None,
         test_transforms: TransformsCompose | None = None,
@@ -185,6 +192,11 @@ def __init__(
                 splits. Only keys listed here reach the batch. All other
                 keys are dropped. See the class docstring for the
                 per-strategy contract.
+            train_collation_map: Additional per-key strategies applied on
+                top of ``collation_map`` for the train dataloader only.
+                Use this for keys that only train transforms produce, so
+                val/test/predict collation does not expect (and warn
+                about) them.
             train_transforms: Transform pipeline applied to training samples.
             val_transforms: Transform pipeline applied to validation samples.
             test_transforms: Transform pipeline applied to test samples.
@@ -197,7 +209,11 @@ def __init__(
         super().__init__()
 
         self.collation_map: dict[str, CollationStrategy] = dict(collation_map or {})
-        if "offset" in self.collation_map:
+        self.train_collation_map: dict[str, CollationStrategy] = {
+            **self.collation_map,
+            **dict(train_collation_map or {}),
+        }
+        if "offset" in self.train_collation_map:
             raise ValueError("'offset' is a reserved collation key generated by concat inputs.")
         # TransformsCompose for each dataset split
         self.train_transforms: TransformsCompose = train_transforms
@@ -306,7 +322,13 @@ def _create_dataloader(self, split: str) -> DataLoader:
         """
         dataset = getattr(self, f"{split}_dataset")
         cfg: DataLoaderConfig = getattr(self, f"{split}_dataloader_cfg")
-        return DataLoader(dataset=dataset, collate_fn=self.collate_fn, **cfg.to_dataloader_kwargs())
+        return DataLoader(
+            dataset=dataset, collate_fn=self._collate_fn_for(split), **cfg.to_dataloader_kwargs()
+        )
+
+    def _collate_fn_for(self, split: str) -> Callable[[list[dict[str, Any]]], dict[str, Any]]:
+        """Return the collate callable matching the split's collation map."""
+        return self.train_collate_fn if split == "train" else self.collate_fn
 
     def train_dataloader(self) -> DataLoader:
         """Create training dataloader."""
@@ -418,7 +440,33 @@ def _apply_index_concat(
         return torch.cat(values, dim=0) + shift
 
     def collate_fn(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
-        """Collate a batch according to ``self.collation_map``.
+        """Collate a val/test/predict batch according to ``self.collation_map``.
+
+        Args:
+            batch: Non-empty list of per-sample dictionaries.
+
+        Returns:
+            Dictionary of collated values keyed by the names declared in
+            ``self.collation_map``. See :meth:`_collate`.
+        """
+        return self._collate(batch, self.collation_map)
+
+    def train_collate_fn(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
+        """Collate a train batch according to ``self.train_collation_map``.
+
+        Args:
+            batch: Non-empty list of per-sample dictionaries.
+
+        Returns:
+            Dictionary of collated values keyed by the names declared in
+            ``self.train_collation_map``. See :meth:`_collate`.
+        """
+        return self._collate(batch, self.train_collation_map)
+
+    def _collate(
+        self, batch: list[dict[str, Any]], collation_map: Mapping[str, CollationStrategy]
+    ) -> dict[str, Any]:
+        """Collate a batch according to the given collation map.
 
         The function dispatches each declared key to its strategy handler,
         derives a cumulative ``offset`` tensor from the first ``"concat"``
@@ -428,10 +476,11 @@ def collate_fn(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
 
         Args:
             batch: Non-empty list of per-sample dictionaries.
+            collation_map: Per-key collation strategies for this split.
 
         Returns:
             Dictionary of collated values keyed by the names declared in
-            ``self.collation_map``. When any ``"concat"`` key is present, an
+            ``collation_map``. When any ``"concat"`` key is present, an
             additional ``"offset"`` key is added with the inclusive
             cumulative count of the primary point cloud space.
 
@@ -447,13 +496,13 @@ def collate_fn(self, batch: list[dict[str, Any]]) -> dict[str, Any]:
         primary_lengths: torch.Tensor | None = None
         deferred_index_concat: list[tuple[str, list[torch.Tensor]]] = []
 
-        for key, strategy in self.collation_map.items():
+        for key, strategy in collation_map.items():
             missing = [i for i, sample in enumerate(batch) if key not in sample]
             if missing:
                 logger.warning(
                     "Key '%s' declared in collation_map but missing from samples %s. "
-                    "Skipping this key during collation. If this comes from deployment/predict, "
-                    "it is expected for training-only annotation keys.",
+                    "Skipping this key during collation. If the key is only produced by "
+                    "train transforms, declare it in train_collation_map instead.",
                     key,
                     missing,
                 )

autoware_ml/datamodule/common/multiview_detection3d.py

@@ -121,7 +121,9 @@ def __init__(
         filter_frames_with_camera_order: bool = True,
         require_image_files: bool = False,
         streaming: bool = False,
+        annotation_status_field: str | None = None,
         collation_map: Mapping[str, CollationStrategy] | None = None,
+        train_collation_map: Mapping[str, CollationStrategy] | None = None,
         train_transforms: TransformsCompose | None = None,
         val_transforms: TransformsCompose | None = None,
         test_transforms: TransformsCompose | None = None,
@@ -150,7 +152,12 @@ def __init__(
                 require it; non-temporal multiview models such as BEVFusion
                 share this datamodule and keep it false so their dataloaders
                 shuffle samples normally.
+            annotation_status_field: Optional annotation field carrying the
+                per-frame annotation-completeness flag consumed by
+                partial-ignore; samples expose it as ``annotation_status``.
             collation_map: Per-key collation strategy applied across all splits.
+            train_collation_map: Additional per-key strategies applied on top
+                of ``collation_map`` for the training split only.
             train_transforms: Transform pipeline applied to training samples.
             val_transforms: Transform pipeline applied to validation samples.
             test_transforms: Transform pipeline applied to test samples.
@@ -162,6 +169,7 @@ def __init__(
         """
         super().__init__(
             collation_map=collation_map,
+            train_collation_map=train_collation_map,
             train_transforms=train_transforms,
             val_transforms=val_transforms,
             test_transforms=test_transforms,
@@ -178,6 +186,7 @@ def __init__(
         self.filter_frames_with_camera_order = filter_frames_with_camera_order
         self.require_image_files = require_image_files
         self.streaming = streaming
+        self.annotation_status_field = annotation_status_field
         self.ann_files = {
             "train": _resolve_ann_file(data_root, train_ann_file),
             "val": _resolve_ann_file(data_root, val_ann_file),
@@ -192,7 +201,7 @@ def _create_dataloader(self, split: str) -> DataLoader:
         return build_streaming_dataloader(
             getattr(self, f"{split}_dataset"),
             getattr(self, f"{split}_dataloader_cfg"),
-            self.collate_fn,
+            self._collate_fn_for(split),
             shuffle_scenes=split == "train",
         )
 
@@ -245,6 +254,7 @@ def __init__(
         name_mapping: Mapping[str, str] | None = None,
         filter_frames_with_camera_order: bool = True,
         require_image_files: bool = False,
+        annotation_status_field: str | None = None,
         dataset_transforms: Any = None,
     ) -> None:
         """Initialize the multiview detection dataset.
@@ -260,6 +270,9 @@ def __init__(
                 image loading always sees a complete camera set.
             require_image_files: Also verify each camera image exists on disk
                 while filtering.
+            annotation_status_field: Optional annotation field carrying the
+                per-frame annotation-completeness flag consumed by
+                partial-ignore; emitted as ``annotation_status``.
             dataset_transforms: Optional transform pipeline.
         """
         super().__init__(dataset_transforms=dataset_transforms)
@@ -268,6 +281,7 @@ def __init__(
         self.camera_order = camera_order
         self.name_mapping = {} if name_mapping is None else dict(name_mapping)
         self.require_image_files = require_image_files
+        self.annotation_status_field = annotation_status_field
         with open(ann_file, "rb") as file:
             data = pickle.load(file)
         data_infos = load_detection_data_infos(data)
@@ -281,7 +295,12 @@ def __init__(
 
     @staticmethod
     def _build_prev_exists(data_infos: list[dict[str, Any]]) -> np.ndarray:
-        """Build stream-continuity flags from adjacent scene tokens."""
+        """Build stream-continuity flags from adjacent scene tokens.
+
+        Camera-order filtering may have dropped mid-scene frames, so a flagged
+        neighbor pair can span a larger time gap; the model consumes the real
+        ``timestamp`` deltas, which absorb this.
+        """
         prev_exists = np.zeros(len(data_infos), dtype=np.float32)
         for index in range(1, len(data_infos)):
             prev_exists[index] = np.float32(
@@ -336,14 +355,26 @@ def __len__(self) -> int:
 
     @staticmethod
     def _build_scene_index_groups(data_infos: list[dict[str, Any]]) -> list[list[int]]:
-        """Group dataset indices by scene, preserving frame order."""
+        """Group dataset indices by scene, requiring scene-contiguous file order.
+
+        ``prev_exists`` is derived from file adjacency, so an annotation file
+        that interleaves scenes would silently reset temporal memory mid-scene;
+        fail loudly instead.
+        """
         groups: dict[str, list[int]] = {}
         for index, sample in enumerate(data_infos):
-            groups.setdefault(sample["scene_token"], []).append(index)
+            token = sample["scene_token"]
+            group = groups.setdefault(token, [])
+            if group and group[-1] != index - 1:
+                raise ValueError(
+                    f"Annotation file interleaves scene '{token}' (frame {index} follows a "
+                    f"gap after frame {group[-1]}); scene frames must be contiguous."
+                )
+            group.append(index)
         return list(groups.values())
 
     def scene_index_groups(self) -> list[list[int]]:
-        """Group dataset indices by scene, preserving frame order.
+        """Return the per-scene dataset index groups.
 
         Returns:
             One list of scene-contiguous dataset indices per scene. A fresh
@@ -465,6 +496,9 @@ def get_data_info(self, index: int) -> dict[str, Any]:
             "scene_token": sample["scene_token"],
             "prev_exists": self.prev_exists[index],
         }
+        if self.annotation_status_field is not None:
+            # Frames without the field count as fully annotated.
+            data_info["annotation_status"] = bool(sample.get(self.annotation_status_field, True))
         ego_pose = _build_ego_pose(sample)
         if ego_pose is not None:
             data_info["ego_pose"] = ego_pose

autoware_ml/datamodule/common/serialization.py

@@ -10,6 +10,7 @@
 
 from __future__ import annotations
 
+import operator
 import pickle
 from collections.abc import Sequence
 from typing import Any
@@ -57,6 +58,7 @@ def __getitem__(self, index: int) -> dict[str, Any]:
         """
         if isinstance(index, slice):
             raise TypeError("SerializedSampleList does not support slicing.")
+        index = operator.index(index)
         length = len(self)
         if index < 0:
             index += length

autoware_ml/datamodule/nuscenes/multiview_detection3d.py

@@ -94,5 +94,6 @@ def _create_dataset(
             name_mapping=self.name_mapping,
             filter_frames_with_camera_order=self.filter_frames_with_camera_order,
             require_image_files=self.require_image_files,
+            annotation_status_field=self.annotation_status_field,
             dataset_transforms=dataset_transforms,
         )

autoware_ml/datamodule/t4dataset/multiview_detection3d.py

@@ -62,5 +62,6 @@ def _create_dataset(
             name_mapping=self.name_mapping,
             filter_frames_with_camera_order=self.filter_frames_with_camera_order,
             require_image_files=self.require_image_files,
+            annotation_status_field=self.annotation_status_field,
             dataset_transforms=dataset_transforms,
         )

autoware_ml/tests/datamodule/test_point_cloud.py

@@ -285,3 +285,41 @@ def test_raises_when_index_concat_has_no_concat_key(self):
         dm = _make_seg_datamodule(collation_map={"inverse": "index_concat"})
         with pytest.raises(ValueError, match="index_concat"):
             dm.collate_fn(batch)
+
+
+class TestTrainCollationMap:
+    def _make_datamodule(self) -> T4Segmentation3DDataModule:
+        return _make_seg_datamodule(
+            collation_map={"coord": "concat"},
+            train_collation_map={"segment": "concat"},
+        )
+
+    def test_train_collate_includes_train_only_keys(self):
+        dm = self._make_datamodule()
+        collated = dm.train_collate_fn(_seg_batch())
+        assert "coord" in collated
+        assert torch.equal(collated["segment"], torch.tensor([0, 1, 2], dtype=torch.long))
+
+    def test_val_collate_ignores_train_only_keys_without_warning(
+        self, caplog: pytest.LogCaptureFixture
+    ):
+        dm = self._make_datamodule()
+        batch = [{"coord": torch.zeros(2, 3)}, {"coord": torch.zeros(3, 3)}]
+
+        collated = dm.collate_fn(batch)
+
+        assert "segment" not in collated
+        assert "missing from samples" not in caplog.text
+
+    def test_collate_fn_for_selects_map_by_split(self):
+        dm = self._make_datamodule()
+        assert dm._collate_fn_for("train") == dm.train_collate_fn
+        for split in ("val", "test", "predict"):
+            assert dm._collate_fn_for(split) == dm.collate_fn
+
+    def test_reserved_offset_key_rejected_in_train_collation_map(self):
+        with pytest.raises(ValueError, match="offset"):
+            _make_seg_datamodule(
+                collation_map={"coord": "concat"},
+                train_collation_map={"offset": "concat"},
+            )

autoware_ml/tests/datamodule/test_serialization.py

@@ -33,6 +33,16 @@ def test_roundtrip_matches_original(self) -> None:
             assert restored["instances"] == original["instances"]
             assert restored["lidar_points"] == original["lidar_points"]
             np.testing.assert_array_equal(restored["array"], original["array"])
+            assert restored["array"].dtype == np.float32
+            # Epoch-second timestamps rely on float64 surviving the round trip.
+            assert isinstance(restored["timestamp"], float)
+            assert restored["timestamp"] == original["timestamp"]
+
+    def test_non_integer_index_is_rejected(self) -> None:
+        serialized = SerializedSampleList(_make_samples(2))
+
+        with pytest.raises(TypeError):
+            serialized[0.0]
 
     def test_negative_index_and_bounds(self) -> None:
         samples = _make_samples(3)