From 76df5cbb8c786bdede6e2c902e9052bd0b56e12a Mon Sep 17 00:00:00 2001 From: nan Date: Wed, 2 Sep 2026 16:00:10 +0800 Subject: [PATCH] Fix temporal BEV alignment for batched tracking Select each sample's previous BEV from the interleaved temporal queue before building attention features. Add numerical coverage for batch alignment and malformed queue lengths. --- .../uniad/modules/temporal_self_attention.py | 21 ++- tests/test_temporal_self_attention.py | 168 ++++++++++++++++++ 2 files changed, 188 insertions(+), 1 deletion(-) create mode 100644 tests/test_temporal_self_attention.py diff --git a/projects/mmdet3d_plugin/uniad/modules/temporal_self_attention.py b/projects/mmdet3d_plugin/uniad/modules/temporal_self_attention.py index f846b4be..0d669c12 100644 --- a/projects/mmdet3d_plugin/uniad/modules/temporal_self_attention.py +++ b/projects/mmdet3d_plugin/uniad/modules/temporal_self_attention.py @@ -21,6 +21,24 @@ '_ext', ['ms_deform_attn_backward', 'ms_deform_attn_forward']) +def _select_previous_bev(value, batch_size, num_bev_queue): + """Select each sample's history BEV from the interleaved queue. + + The encoder packs temporal values as ``[prev_0, curr_0, prev_1, + curr_1, ...]``. Flattened indexing with ``value[:batch_size]`` only + selects the correct history for the first sample when ``batch_size > 1``. + Reshaping the queue restores the sample/queue axes before selecting the + history slot. + """ + if value.shape[0] != batch_size * num_bev_queue: + raise ValueError( + 'Expected value batch dimension to equal ' + f'batch_size * num_bev_queue ({batch_size} * {num_bev_queue}), ' + f'but got {value.shape[0]}' + ) + return value.reshape(batch_size, num_bev_queue, *value.shape[1:])[:, 0] + + @ATTENTION.register_module() class TemporalSelfAttention(BaseModule): """An attention module used in BEVFormer based on Deformable-Detr. @@ -191,7 +209,8 @@ def forward(self, assert (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == num_value assert self.num_bev_queue == 2 - query = torch.cat([value[:bs], query], -1) + query = torch.cat( + [_select_previous_bev(value, bs, self.num_bev_queue), query], -1) value = self.value_proj(value) if key_padding_mask is not None: diff --git a/tests/test_temporal_self_attention.py b/tests/test_temporal_self_attention.py new file mode 100644 index 00000000..76789040 --- /dev/null +++ b/tests/test_temporal_self_attention.py @@ -0,0 +1,168 @@ +import importlib.util +import sys +import types +from pathlib import Path + +import pytest +import torch +from torch import nn + + +def _stub(monkeypatch, name, package=False, **attributes): + module = types.ModuleType(name) + if package: + module.__path__ = [] + for key, value in attributes.items(): + setattr(module, key, value) + monkeypatch.setitem(sys.modules, name, module) + return module + + +def _load_temporal_attention(monkeypatch): + """Load the module while replacing its optional simulator dependencies.""" + package_name = "_uniad_temporal_attention_test" + source_dir = Path(__file__).resolve().parents[1] / "projects/mmdet3d_plugin/uniad/modules" + package = _stub(monkeypatch, package_name, package=True) + package.__path__ = [str(source_dir)] + _stub( + monkeypatch, + f"{package_name}.multi_scale_deformable_attn_function", + MultiScaleDeformableAttnFunction_fp32=object, + ) + + _stub(monkeypatch, "mmcv", package=True) + _stub(monkeypatch, "mmcv.ops", package=True) + _stub( + monkeypatch, + "mmcv.ops.multi_scale_deform_attn", + multi_scale_deformable_attn_pytorch=lambda *args, **kwargs: None, + ) + _stub( + monkeypatch, + "mmcv.cnn", + xavier_init=lambda *args, **kwargs: None, + constant_init=lambda *args, **kwargs: None, + ) + _stub(monkeypatch, "mmcv.cnn.bricks", package=True) + + class Registry: + def register_module(self): + return lambda cls: cls + + _stub(monkeypatch, "mmcv.cnn.bricks.registry", ATTENTION=Registry()) + _stub(monkeypatch, "mmcv.runner", package=True) + + class BaseModule(nn.Module): + def __init__(self, init_cfg=None): + super().__init__() + + _stub( + monkeypatch, + "mmcv.runner.base_module", + BaseModule=BaseModule, + ModuleList=nn.ModuleList, + Sequential=nn.Sequential, + ) + _stub(monkeypatch, "mmcv.utils", package=True) + ext_loader = _stub( + monkeypatch, + "mmcv.utils.ext_loader", + load_ext=lambda *args, **kwargs: types.SimpleNamespace(), + ) + utils = sys.modules["mmcv.utils"] + utils.ConfigDict = dict + utils.build_from_cfg = lambda *args, **kwargs: None + utils.deprecated_api_warning = lambda *args, **kwargs: None + utils.to_2tuple = lambda value: (value, value) + utils.ext_loader = ext_loader + + module_name = f"{package_name}.temporal_self_attention" + spec = importlib.util.spec_from_file_location( + module_name, source_dir / "temporal_self_attention.py" + ) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, module_name, module) + spec.loader.exec_module(module) + return module + + +def test_previous_bev_selection_preserves_sample_alignment(monkeypatch): + temporal_attention = _load_temporal_attention(monkeypatch) + value = torch.tensor( + [ + [[10.0, 10.1]], # sample 0, previous frame + [[20.0, 20.1]], # sample 0, current frame + [[30.0, 30.1]], # sample 1, previous frame + [[40.0, 40.1]], # sample 1, current frame + ] + ) + + previous = temporal_attention._select_previous_bev(value, 2, 2) + + torch.testing.assert_close(previous, value[[0, 2]]) + assert not torch.equal(previous, value[:2]) + + +def test_temporal_attention_uses_each_sample_history(monkeypatch): + temporal_attention = _load_temporal_attention(monkeypatch) + attention = temporal_attention.TemporalSelfAttention( + embed_dims=2, num_heads=1, num_levels=1, num_points=1, dropout=0.0 + ) + + class CaptureProjection(nn.Module): + def __init__(self, output_size, capture=False): + super().__init__() + self.output_size = output_size + self.capture = capture + + def forward(self, input_tensor): + if self.capture: + self.input = input_tensor.detach().clone() + return input_tensor.new_zeros( + input_tensor.shape[0], input_tensor.shape[1], self.output_size + ) + + offsets = CaptureProjection(4, capture=True) + attention.sampling_offsets = offsets + attention.attention_weights = CaptureProjection(2) + attention.value_proj = nn.Identity() + attention.output_proj = nn.Identity() + attention.dropout = nn.Identity() + monkeypatch.setattr( + temporal_attention, + "multi_scale_deformable_attn_pytorch", + lambda value, spatial_shapes, sampling_locations, attention_weights: value.new_zeros( + value.shape[0], sampling_locations.shape[1], value.shape[-2] * value.shape[-1] + ), + ) + + attention( + query=torch.tensor([[[100.0, 101.0]], [[200.0, 201.0]]]), + value=torch.tensor( + [ + [[10.0, 11.0]], # sample 0, previous frame + [[20.0, 21.0]], # sample 0, current frame + [[30.0, 31.0]], # sample 1, previous frame + [[40.0, 41.0]], # sample 1, current frame + ] + ), + reference_points=torch.zeros(4, 1, 1, 2), + spatial_shapes=torch.tensor([[1, 1]]), + level_start_index=torch.tensor([0]), + ) + + torch.testing.assert_close( + offsets.input, + torch.tensor( + [ + [[10.0, 11.0, 100.0, 101.0]], + [[30.0, 31.0, 200.0, 201.0]], + ] + ), + ) + + +def test_previous_bev_selection_rejects_incomplete_queue(monkeypatch): + temporal_attention = _load_temporal_attention(monkeypatch) + with pytest.raises(ValueError, match=r"batch_size \* num_bev_queue"): + temporal_attention._select_previous_bev(torch.zeros(3, 1, 4), 2, 2)