Merge pull request #61 from apgeorg/test/split-test-suite

test: split test_auto_e2e into per-package test modules

Author: riita10069

Model/tests/test_auto_e2e.py

@@ -0,0 +1,110 @@
+import torch
+import sys
+sys.path.append('..')
+
+from model_components.backbone import Backbone
+
+
+class _StubBackboneWithFeatureInfo(torch.nn.Module):
+    """Channels-first backbone exposing timm-style feature_info."""
+
+    def __init__(self):
+        super().__init__()
+        self.stage0 = torch.nn.Conv2d(3, 32, 3, stride=2, padding=1)
+        self.stage1 = torch.nn.Conv2d(32, 48, 3, stride=2, padding=1)
+        self.stage2 = torch.nn.Conv2d(48, 64, 3, stride=2, padding=1)
+        self.feature_info = [{"num_chs": 32}, {"num_chs": 48}, {"num_chs": 64}]
+
+    def forward(self, x):
+        s0 = self.stage0(x)
+        s1 = self.stage1(s0)
+        s2 = self.stage2(s1)
+        return [s0, s1, s2]
+
+
+class _StubBackboneNoFeatureInfo(torch.nn.Module):
+    """Channels-first backbone with NO feature_info (probe fallback path)."""
+
+    def __init__(self):
+        super().__init__()
+        self.stage0 = torch.nn.Conv2d(3, 24, 3, stride=2, padding=1)
+        self.stage1 = torch.nn.Conv2d(24, 56, 3, stride=2, padding=1)
+        self.stage2 = torch.nn.Conv2d(56, 112, 3, stride=2, padding=1)
+
+    def forward(self, x):
+        s0 = self.stage0(x)
+        s1 = self.stage1(s0)
+        s2 = self.stage2(s1)
+        return [s0, s1, s2]
+
+
+class _StubBackboneSwinLike(torch.nn.Module):
+    """Channels-last backbone (B, H, W, C) — exercises permute branch."""
+
+    def __init__(self):
+        super().__init__()
+        self.stage0 = torch.nn.Conv2d(3, 32, 3, stride=2, padding=1)
+        self.stage1 = torch.nn.Conv2d(32, 48, 3, stride=2, padding=1)
+        self.feature_info = [{"num_chs": 32}, {"num_chs": 48}]
+
+    def forward(self, x):
+        s0_cf = self.stage0(x)                                  # [B, 32, H, W]
+        s1_cf = self.stage1(s0_cf)                              # [B, 48, H, W]
+        s0 = s0_cf.permute(0, 2, 3, 1).contiguous()             # [B, H, W, 32]
+        s1 = s1_cf.permute(0, 2, 3, 1).contiguous()             # [B, H, W, 48]
+        return [s0, s1]
+
+
+class TestBackboneChannelDiscovery:
+    """Cover the backbone_channels discovery + layout-detection in Backbone."""
+
+    def _make_backbone(self, monkeypatch, stub_module):
+        # Patch the registry call so build_backbone returns our stub.
+        monkeypatch.setattr(
+            "model_components.backbone.build_backbone",
+            lambda *a, **kw: stub_module,
+        )
+        return Backbone(backbone="stub", is_pretrained=False)
+
+    def test_feature_info_path_sums_channels(self, monkeypatch):
+        bb = self._make_backbone(monkeypatch, _StubBackboneWithFeatureInfo())
+        assert bb.backbone_channels == 32 + 48 + 64
+
+    def test_probe_fallback_when_feature_info_missing(self, monkeypatch):
+        bb = self._make_backbone(monkeypatch, _StubBackboneNoFeatureInfo())
+        # No feature_info — channels recovered via probing.
+        assert bb.backbone_channels == 24 + 56 + 112
+
+    def test_feature_info_channels_match_forward_output(self, monkeypatch, device):
+        """sum(feature_info channels) must equal the actual concat-channel dim
+        of the forward output."""
+        bb = self._make_backbone(monkeypatch, _StubBackboneWithFeatureInfo()).to(device)
+        x = torch.randn(2, 3, 32, 32, device=device)
+        feats = bb(x)
+        total_c = sum(f.shape[1] for f in feats)
+        assert total_c == bb.backbone_channels
+
+    def test_probe_channels_match_forward_output(self, monkeypatch, device):
+        bb = self._make_backbone(monkeypatch, _StubBackboneNoFeatureInfo()).to(device)
+        x = torch.randn(2, 3, 32, 32, device=device)
+        feats = bb(x)
+        total_c = sum(f.shape[1] for f in feats)
+        assert total_c == bb.backbone_channels
+
+    def test_channels_last_backbone_is_permuted(self, monkeypatch, device):
+        """Channels-last (B, H, W, C) output must be permuted to (B, C, H, W)
+        based on tensor shape, NOT on the backbone name."""
+        bb = self._make_backbone(monkeypatch, _StubBackboneSwinLike()).to(device)
+        x = torch.randn(2, 3, 32, 32, device=device)
+        feats = bb(x)
+        # After Backbone.forward, every feature must be channels-first with the
+        # expected channel count at dim 1.
+        assert feats[0].shape[1] == 32
+        assert feats[1].shape[1] == 48
+
+    def test_channels_first_backbone_not_permuted(self, monkeypatch, device):
+        bb = self._make_backbone(monkeypatch, _StubBackboneWithFeatureInfo()).to(device)
+        x = torch.randn(2, 3, 32, 32, device=device)
+        feats = bb(x)
+        for f, expected in zip(feats, [32, 48, 64]):
+            assert f.shape[1] == expected

Model/tests/test_backbone.py

@@ -0,0 +1,58 @@
+import torch
+import sys
+sys.path.append('..')
+
+from model_components.feature_fusion import FeatureFusion
+
+
+class TestFeatureFusionComponent:
+    def test_output_shape(self, device):
+        fusion = FeatureFusion(num_views=8, fusion_mode="concat").to(device)
+        features = [
+            torch.randn(16, 96, 64, 64, device=device),
+            torch.randn(16, 192, 32, 32, device=device),
+            torch.randn(16, 384, 16, 16, device=device),
+            torch.randn(16, 768, 8, 8, device=device),
+        ]
+        out = fusion(features, B=2, V=8)
+        assert out.shape == (2, 256, 8, 8)
+
+    def test_view_reduction_changes_output(self, device):
+        """Verify that view_reduce is not identity (actually mixes views)."""
+        fusion = FeatureFusion(num_views=8, fusion_mode="concat").to(device)
+        fusion.eval()
+
+        features_a = [
+            torch.randn(8, 96, 64, 64, device=device),
+            torch.randn(8, 192, 32, 32, device=device),
+            torch.randn(8, 384, 16, 16, device=device),
+            torch.randn(8, 768, 8, 8, device=device),
+        ]
+        out_a = fusion(features_a, B=1, V=8)
+
+        features_b = [f.clone() for f in features_a]
+        features_b[0][3] = torch.randn_like(features_b[0][3])
+        out_b = fusion(features_b, B=1, V=8)
+
+        assert not torch.allclose(out_a, out_b, atol=1e-5)
+
+
+class TestFeatureFusionWithSwinChannels:
+    def test_dynamic_backbone_channels_with_swin_sizes(self, device):
+        """FeatureFusion should accept Swin's per-stage channels (96, 192, 384, 768)
+        at their natural spatial resolutions and produce the expected fused shape."""
+        backbone_channels = 96 + 192 + 384 + 768  # 1440
+        fusion = FeatureFusion(
+            num_views=8, backbone_channels=backbone_channels, fusion_mode="concat",
+        ).to(device)
+
+        # Per-stage Swin spatial dims for a 256x256 input
+        features = [
+            torch.randn(16, 96, 64, 64, device=device),
+            torch.randn(16, 192, 32, 32, device=device),
+            torch.randn(16, 384, 16, 16, device=device),
+            torch.randn(16, 768, 8, 8, device=device),
+        ]
+        out = fusion(features, B=2, V=8)
+        assert out.shape == (2, 256, 8, 8)
+        assert torch.isfinite(out).all()

Model/tests/test_feature_fusion.py

@@ -0,0 +1,61 @@
+import torch
+import sys
+sys.path.append('..')
+
+from model_components.future_state import FutureState
+
+
+class TestFutureStateComponent:
+    def test_accepts_ego_hidden(self, device):
+        future = FutureState(embed_dim=256, ego_hidden_dim=256).to(device)
+        feats = torch.randn(2, 256, 8, 8, device=device)
+        ego_hidden = torch.randn(2, 256, device=device)
+        out = future(feats, ego_hidden)
+        assert len(out) == 4
+        for f in out:
+            assert f.shape == (2, 256, 8, 8)
+
+    def test_ego_hidden_influences_output(self, device):
+        future = FutureState(embed_dim=256, ego_hidden_dim=256).to(device)
+        future.eval()
+        feats = torch.randn(1, 256, 8, 8, device=device)
+
+        out_a = future(feats, torch.randn(1, 256, device=device))
+        out_b = future(feats, torch.randn(1, 256, device=device))
+
+        assert not torch.allclose(out_a[0], out_b[0], atol=1e-5), \
+            "ego_hidden should influence future predictions"
+
+
+class TestFutureStateChunkSplit:
+    def test_four_outputs_are_distinct(self, device):
+        """torch.chunk must split along channels, not return 4 views of the same data."""
+        torch.manual_seed(0)
+        future = FutureState(embed_dim=256, ego_hidden_dim=256).to(device)
+        future.eval()
+        feats = torch.randn(2, 256, 8, 8, device=device)
+        ego_hidden = torch.randn(2, 256, device=device)
+
+        out = future(feats, ego_hidden)
+        assert len(out) == 4
+        for i in range(4):
+            for j in range(i + 1, 4):
+                assert not torch.allclose(out[i], out[j], atol=1e-5), \
+                    f"FutureState outputs {i} and {j} are identical — chunk is broken"
+
+    def test_ego_hidden_changes_all_four_outputs(self, device):
+        """Different ego_hidden must shift every one of the 4 future predictions."""
+        torch.manual_seed(0)
+        future = FutureState(embed_dim=256, ego_hidden_dim=256).to(device)
+        future.eval()
+        feats = torch.randn(1, 256, 8, 8, device=device)
+
+        ego_a = torch.randn(1, 256, device=device)
+        ego_b = torch.randn(1, 256, device=device)
+
+        out_a = future(feats, ego_a)
+        out_b = future(feats, ego_b)
+
+        for i in range(4):
+            assert not torch.allclose(out_a[i], out_b[i], atol=1e-5), \
+                f"Future output {i} did not change when ego_hidden changed"

Model/tests/test_future_state.py

@@ -0,0 +1,95 @@
+import pytest
+import torch
+import sys
+sys.path.append('..')
+
+
+def make_inputs(batch_size, num_views, device, include_camera_params=False):
+    visual = torch.randn(batch_size, num_views, 3, 256, 256, device=device)
+    map_input = torch.randn(batch_size, 3, 256, 256, device=device)
+    visual_history = torch.randn(batch_size, 896, device=device)
+    egomotion = torch.randn(batch_size, 256, device=device)
+    if include_camera_params:
+        camera_params = torch.randn(batch_size, num_views, 3, 4, device=device)
+        return visual, map_input, visual_history, egomotion, camera_params
+    return visual, map_input, visual_history, egomotion
+
+
+# ---------------------------------------------------------------------------
+# Integration tests — full backbone (slow, marked for separate CI tier)
+# ---------------------------------------------------------------------------
+
+@pytest.mark.integration
+class TestFullBackboneIntegration:
+    """End-to-end tests with the real pretrained backbone.
+
+    These verify that the full pipeline (backbone → fusion → planner → future)
+    produces correct shapes and numerically stable outputs. Run separately
+    from unit tests via: pytest -m integration
+    """
+
+    def test_full_forward_pass(self, full_model, device):
+        """Smoke test: full model forward produces expected output shapes."""
+        visual, map_input, vis_hist, ego = make_inputs(1, 7, device)
+        target = torch.randn(1, 128, device=device)
+        loss, ego_hidden, future = full_model(
+             visual, map_input, vis_hist, ego,
+            mode="train", trajectory_target=target)
+
+        assert loss.dim() == 0
+        assert ego_hidden.shape == (1, 256)
+        assert len(future) == 4
+        for f in future:
+            assert f.shape == (1, 256, 8, 8)
+
+        traj, _, _ = full_model(visual, vis_hist, ego, mode="infer")
+        assert traj.shape == (1, 128)
+
+    def test_full_forward_no_nan(self, full_model, device):
+        """Full pipeline must not produce NaN with real backbone weights."""
+        visual, map_input, vis_hist, ego = make_inputs(2, 7, device)
+        target = torch.randn(2, 128, device=device)
+        loss, ego_hidden, future = full_model(visual, map_input, vis_hist, ego,
+                                              mode="train", trajectory_target=target)
+
+        assert not torch.isnan(loss)
+        assert not torch.isnan(ego_hidden).any()
+        for f in future:
+            assert not torch.isnan(f).any()
+
+
+@pytest.mark.integration
+class TestResNet50Backbone:
+    """Exercises the dynamic backbone_channels computation on a backbone
+    whose feature_info shape differs from Swin (5 stages of channels
+    64/256/512/1024/2048 vs Swin's 4 stages of 96/192/384/768)."""
+
+    def test_resnet50_forward_pass(self, device):
+        from model_components.auto_e2e import AutoE2E
+        try:
+            model = AutoE2E(
+                backbone="res_net_50", num_views=7, fusion_mode="concat",
+                is_pretrained=False,
+            ).to(device)
+        except (FileNotFoundError, OSError) as e:
+            pytest.skip(f"Backbone construction failed: {e}")
+
+        # Dynamic backbone_channels = sum of all 5 ResNet50 stages = 3904
+        assert model.Backbone.backbone_channels == 64 + 256 + 512 + 1024 + 2048
+
+        visual, map_input, vis_hist, ego = make_inputs(1, 7, device)
+        target = torch.randn(1, 128, device=device)
+        loss, ego_hidden, future = model(
+            visual, map_input, vis_hist, ego, mode="train", trajectory_target=target)
+
+        assert loss.dim() == 0
+        assert ego_hidden.shape == (1, 256)
+        assert len(future) == 4
+        for f in future:
+            assert f.shape == (1, 256, 8, 8)
+        assert torch.isfinite(loss)
+        assert torch.isfinite(ego_hidden).all()
+
+        traj, _, _ = model(visual, vis_hist, ego, mode="infer")
+        assert traj.shape == (1, 128)
+        assert torch.isfinite(traj).all()