Add DeMemWM noise-route pruning diagnostic

.exp_artifact/dememwm_noise_route_inference_pruning/diagnostic.py

@@ -0,0 +1,236 @@
+from pathlib import Path
+from types import SimpleNamespace
+import sys
+
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from algorithms.dememwm.df_video import (  # noqa: E402
+    _DEMEMWM_STREAM_KEYS,
+    _apply_memory_route_masks,
+    _pack_active_inference_memory_streams,
+)
+
+
+ROUTE_CFG = {"noise_route": {"anchor": "high", "dynamic": "low", "revisit": "all"}}
+DIFFUSION = SimpleNamespace(timesteps=100, sampling_timesteps=4)
+STREAM_LENGTHS = {"anchor": 2, "dynamic": 3, "revisit": 2}
+STREAM_VALUES = {"anchor": 20.0, "dynamic": 30.0, "revisit": 40.0}
+STREAM_NOISE_BASE = {"anchor": 200, "dynamic": 300, "revisit": 400}
+BATCH_SIZE = 2
+TARGET_LENGTH = 3
+
+
+def _make_inputs(stream_masks):
+    target_latents = torch.arange(TARGET_LENGTH * BATCH_SIZE * 2, dtype=torch.float32).reshape(TARGET_LENGTH, BATCH_SIZE, 2)
+    target_conditions = torch.arange(TARGET_LENGTH * BATCH_SIZE * 4, dtype=torch.float32).reshape(TARGET_LENGTH, BATCH_SIZE, 4)
+    target_poses = torch.arange(TARGET_LENGTH * BATCH_SIZE * 3, dtype=torch.float32).reshape(TARGET_LENGTH, BATCH_SIZE, 3)
+    target_frame_indices = torch.arange(TARGET_LENGTH, dtype=torch.long)[:, None].expand(TARGET_LENGTH, BATCH_SIZE).clone()
+    target_mask = torch.tensor([[True, False, True], [False, True, True]], dtype=torch.bool)
+
+    stream_latents_by_key = {}
+    stream_poses_by_key = {}
+    stream_frame_indices_by_key = {}
+    memory_noise_levels_by_key = {}
+    for key in _DEMEMWM_STREAM_KEYS:
+        length = STREAM_LENGTHS[key]
+        stream_latents_by_key[key] = torch.full((length, BATCH_SIZE, 2), STREAM_VALUES[key])
+        stream_poses_by_key[key] = torch.full((length, BATCH_SIZE, 3), STREAM_VALUES[key] + 0.5)
+        stream_frame_indices_by_key[key] = (
+            torch.arange(length, dtype=torch.long)[:, None].expand(length, BATCH_SIZE).clone()
+            + int(STREAM_VALUES[key])
+        )
+        memory_noise_levels_by_key[key] = (
+            torch.arange(length * BATCH_SIZE, dtype=torch.long).reshape(length, BATCH_SIZE)
+            + STREAM_NOISE_BASE[key]
+        )
+
+    frame_memory_masks = {"target": target_mask, **stream_masks}
+    return {
+        "target_latents": target_latents,
+        "target_conditions": target_conditions,
+        "target_poses": target_poses,
+        "target_frame_indices": target_frame_indices,
+        "target_mask": target_mask,
+        "stream_latents_by_key": stream_latents_by_key,
+        "stream_poses_by_key": stream_poses_by_key,
+        "stream_frame_indices_by_key": stream_frame_indices_by_key,
+        "frame_memory_masks": frame_memory_masks,
+        "memory_noise_levels_by_key": memory_noise_levels_by_key,
+    }
+
+
+def _sampler_index_noise(index):
+    real_steps = torch.linspace(-1, DIFFUSION.timesteps - 1, steps=DIFFUSION.sampling_timesteps + 1).long()
+    return int(real_steps[index].item())
+
+
+def _target_proxy(packed_latents, frame_memory_segments, frame_memory_masks):
+    target_len = frame_memory_segments["target"]
+    target = packed_latents[:target_len].clone()
+    cursor = target_len
+    contribution = target.new_zeros((BATCH_SIZE, packed_latents.shape[-1]))
+    for key_index, key in enumerate(_DEMEMWM_STREAM_KEYS, start=1):
+        length = frame_memory_segments[key]
+        segment = packed_latents[cursor : cursor + length]
+        mask = frame_memory_masks[key]
+        cursor += length
+        if length == 0 or mask.numel() == 0:
+            continue
+        visible = mask.T.to(dtype=segment.dtype).unsqueeze(-1)
+        contribution = contribution + float(key_index) * (segment * visible).sum(dim=0)
+    return target + contribution.unsqueeze(0)
+
+
+def _full_proxy_inputs(inputs, routed_masks):
+    packed_latents = torch.cat(
+        [inputs["target_latents"], *[inputs["stream_latents_by_key"][key] for key in _DEMEMWM_STREAM_KEYS]],
+        dim=0,
+    )
+    segments = {"target": TARGET_LENGTH, **{key: STREAM_LENGTHS[key] for key in _DEMEMWM_STREAM_KEYS}}
+    return packed_latents, segments, {"target": inputs["target_mask"], **routed_masks}
+
+
+def _expected_noise(memory_noise_levels_by_key, active_streams):
+    if not active_streams:
+        return torch.zeros((0, BATCH_SIZE), dtype=torch.long)
+    return torch.cat([memory_noise_levels_by_key[key] for key in active_streams], dim=0)
+
+
+def _check(condition, message):
+    if not condition:
+        raise AssertionError(message)
+
+
+def _run_case(name, sampler_index, stream_masks, expected_active, expected_pruned):
+    inputs = _make_inputs(stream_masks)
+    query_noise_levels = torch.full((TARGET_LENGTH, BATCH_SIZE), sampler_index, dtype=torch.long)
+    routed_masks = _apply_memory_route_masks(
+        inputs["frame_memory_masks"],
+        query_noise_levels,
+        ROUTE_CFG,
+        DIFFUSION,
+        mode="validation",
+    )
+    _check(torch.equal(routed_masks["target"], inputs["target_mask"]), f"{name}: target mask was routed")
+
+    (
+        active_packed_latents,
+        active_packed_conditions,
+        active_frame_memory_pose,
+        active_frame_indices,
+        active_segments,
+        active_masks,
+        active_memory_noise_levels,
+        active_streams,
+        pruned_streams,
+    ) = _pack_active_inference_memory_streams(
+        inputs["target_latents"],
+        inputs["target_conditions"],
+        inputs["target_poses"],
+        inputs["target_frame_indices"],
+        inputs["target_mask"],
+        inputs["stream_latents_by_key"],
+        inputs["stream_poses_by_key"],
+        inputs["stream_frame_indices_by_key"],
+        routed_masks,
+        inputs["memory_noise_levels_by_key"],
+    )
+
+    full_segments = {"target": TARGET_LENGTH, **{key: STREAM_LENGTHS[key] for key in _DEMEMWM_STREAM_KEYS}}
+    full_memory_tokens = sum(full_segments[key] for key in _DEMEMWM_STREAM_KEYS)
+    pruned_memory_tokens = sum(active_segments[key] for key in _DEMEMWM_STREAM_KEYS)
+
+    _check(active_streams == expected_active, f"{name}: active streams {active_streams} != {expected_active}")
+    _check(pruned_streams == expected_pruned, f"{name}: pruned streams {pruned_streams} != {expected_pruned}")
+    _check(set(active_segments) == {"target", *_DEMEMWM_STREAM_KEYS}, f"{name}: missing segment keys")
+    _check(active_segments["target"] == TARGET_LENGTH, f"{name}: target segment length changed")
+    for key in expected_pruned:
+        _check(active_segments[key] == 0, f"{name}: pruned {key} has nonzero segment")
+        _check(active_masks[key].shape == (BATCH_SIZE, 0), f"{name}: pruned {key} mask is not zero-length")
+    for key in expected_active:
+        _check(active_segments[key] == STREAM_LENGTHS[key], f"{name}: active {key} segment length changed")
+        _check(torch.equal(active_masks[key], routed_masks[key]), f"{name}: active {key} mask changed")
+
+    expected_total_len = TARGET_LENGTH + pruned_memory_tokens
+    _check(active_packed_latents.shape == (expected_total_len, BATCH_SIZE, 2), f"{name}: packed latent shape mismatch")
+    _check(active_packed_conditions.shape == (expected_total_len, BATCH_SIZE, 4), f"{name}: packed condition shape mismatch")
+    _check(active_frame_memory_pose.shape == (expected_total_len, BATCH_SIZE, 3), f"{name}: packed pose shape mismatch")
+    _check(active_frame_indices.shape == (expected_total_len, BATCH_SIZE), f"{name}: packed frame index shape mismatch")
+    _check(torch.equal(active_packed_conditions[:TARGET_LENGTH], inputs["target_conditions"]), f"{name}: target conditions changed")
+    _check(torch.equal(active_packed_conditions[TARGET_LENGTH:], torch.zeros_like(active_packed_conditions[TARGET_LENGTH:])), f"{name}: memory conditions are not zero")
+
+    cursor = TARGET_LENGTH
+    for key in expected_active:
+        length = STREAM_LENGTHS[key]
+        expected_chunk = torch.full((length, BATCH_SIZE, 2), STREAM_VALUES[key])
+        _check(torch.equal(active_packed_latents[cursor : cursor + length], expected_chunk), f"{name}: canonical packed order broke at {key}")
+        cursor += length
+
+    expected_noise = _expected_noise(inputs["memory_noise_levels_by_key"], expected_active)
+    _check(torch.equal(active_memory_noise_levels, expected_noise), f"{name}: memory-noise concatenation order mismatch")
+    from_target = torch.full((TARGET_LENGTH, BATCH_SIZE), sampler_index, dtype=torch.long)
+    _check(
+        torch.cat([from_target, active_memory_noise_levels], dim=0).shape[0] == active_packed_latents.shape[0],
+        f"{name}: denoising noise-level shape contract mismatch",
+    )
+
+    full_latents, full_proxy_segments, full_proxy_masks = _full_proxy_inputs(inputs, routed_masks)
+    full_output = _target_proxy(full_latents, full_proxy_segments, full_proxy_masks)
+    pruned_output = _target_proxy(active_packed_latents, active_segments, active_masks)
+    target_max_abs_diff = (full_output - pruned_output).abs().max().item()
+    _check(target_max_abs_diff == 0.0, f"{name}: proxy target output changed by {target_max_abs_diff}")
+
+    print(f"case={name}")
+    print(f"  route=anchor:high,dynamic:low,revisit:all")
+    print(f"  denoising_step={sampler_index} noise_level={_sampler_index_noise(sampler_index)}")
+    print(f"  active_streams={active_streams}")
+    print(f"  pruned_streams={pruned_streams}")
+    print(f"  full_memory_token_count={full_memory_tokens}")
+    print(f"  pruned_memory_token_count={pruned_memory_tokens}")
+    print(f"  segments_before={full_segments}")
+    print(f"  segments_after={active_segments}")
+    print(f"  target_proxy_max_abs_diff={target_max_abs_diff:.6f}")
+
+
+def main():
+    all_valid_masks = {
+        "anchor": torch.tensor([[True, True], [True, False]], dtype=torch.bool),
+        "dynamic": torch.tensor([[True, True, False], [False, True, True]], dtype=torch.bool),
+        "revisit": torch.tensor([[True, False], [True, True]], dtype=torch.bool),
+    }
+    empty_anchor_masks = {
+        **all_valid_masks,
+        "anchor": torch.zeros((BATCH_SIZE, STREAM_LENGTHS["anchor"]), dtype=torch.bool),
+    }
+
+    _run_case(
+        "high_noise_all_valid",
+        sampler_index=4,
+        stream_masks=all_valid_masks,
+        expected_active=["anchor", "revisit"],
+        expected_pruned=["dynamic"],
+    )
+    _run_case(
+        "low_noise_all_valid",
+        sampler_index=1,
+        stream_masks=all_valid_masks,
+        expected_active=["dynamic", "revisit"],
+        expected_pruned=["anchor"],
+    )
+    _run_case(
+        "high_noise_empty_anchor",
+        sampler_index=4,
+        stream_masks=empty_anchor_masks,
+        expected_active=["revisit"],
+        expected_pruned=["anchor", "dynamic"],
+    )
+    print("diagnostic=ok")
+
+
+if __name__ == "__main__":
+    main()