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)