Upload test_all.py with huggingface_hub

test_all.py

@@ -0,0 +1,468 @@
+"""
+Comprehensive test suite for ViL Tracker.
+
+13 tests covering all components:
+1. mLSTM Cell (LinearHeadwiseExpand correctness + param count)
+2. mLSTM Block (full block with MLP)
+3. TMoE MLP
+4. Backbone (standard, small depth)
+5. Backbone (with TMoE, medium depth)
+6. Prediction Heads
+7. FiLM Temporal Modulation
+8. Full Tracker (small depth for speed)
+9. Loss Functions
+10. Kalman Filter
+11. Dataset (synthetic)
+12. Training Step (mini forward + backward)
+13. Model Summary (FULL depth=24, constraint check)
+"""
+
+import sys
+import time
+import torch
+import numpy as np
+
+torch.manual_seed(42)
+np.random.seed(42)
+
+PASS = 0
+FAIL = 0
+
+def test(name, fn):
+    global PASS, FAIL
+    print(f"\nTest {PASS + FAIL + 1}: {name}...", flush=True)
+    try:
+        fn()
+        PASS += 1
+        print(f"  ✅ PASSED")
+    except Exception as e:
+        FAIL += 1
+        print(f"  ❌ FAILED: {e}")
+        import traceback
+        traceback.print_exc()
+
+
+def count_params(model):
+    return sum(p.numel() for p in model.parameters())
+
+
+# ============================================================
+# Test 1: mLSTM Cell
+# ============================================================
+def test_mlstm_cell():
+    from vil_tracker.models.mlstm import mLSTMCell, LinearHeadwiseExpand
+    
+    # Test LinearHeadwiseExpand
+    lhe = LinearHeadwiseExpand(768, num_heads=192, bias=False)
+    lhe_params = count_params(lhe)
+    assert lhe_params == 192 * 4 * 4, f"LHE params: {lhe_params} != {192*4*4}"
+    
+    x = torch.randn(2, 10, 768)
+    y = lhe(x)
+    assert y.shape == (2, 10, 768), f"LHE output shape: {y.shape}"
+    
+    # Test full mLSTM cell
+    cell = mLSTMCell(dim=384, proj_factor=2.0, qkv_proj_blocksize=4, num_heads=4)
+    cell_params = count_params(cell)
+    print(f"  mLSTMCell params: {cell_params:,} ({cell_params/1e6:.3f}M)")
+    
+    # Should be ~920K, not 2.66M
+    assert cell_params < 1_000_000, f"Cell has {cell_params:,} params (should be <1M)"
+    assert cell_params > 800_000, f"Cell has {cell_params:,} params (should be >800K)"
+    
+    x = torch.randn(2, 20, 384)
+    y = cell(x)
+    assert y.shape == (2, 20, 384), f"Cell output shape: {y.shape}"
+    
+    # Test reverse mode
+    y_rev = cell(x, reverse=True)
+    assert y_rev.shape == (2, 20, 384), f"Reverse output shape: {y_rev.shape}"
+    # Forward and reverse should produce different results
+    assert not torch.allclose(y, y_rev, atol=1e-3), "Forward and reverse should differ"
+
+test("mLSTM Cell (LinearHeadwiseExpand)", test_mlstm_cell)
+
+
+# ============================================================
+# Test 2: mLSTM Block
+# ============================================================
+def test_mlstm_block():
+    from vil_tracker.models.mlstm import mLSTMBlock
+    
+    block = mLSTMBlock(dim=384, proj_factor=2.0, qkv_proj_blocksize=4,
+                       num_heads=4, mlp_ratio=4.0)
+    params = count_params(block)
+    print(f"  mLSTMBlock params: {params:,} ({params/1e6:.3f}M)")
+    
+    x = torch.randn(2, 20, 384)
+    y = block(x)
+    assert y.shape == (2, 20, 384), f"Block output shape: {y.shape}"
+    
+    # Residual connection: output should be close-ish to input at init
+    diff = (y - x).abs().mean().item()
+    print(f"  Residual diff from input: {diff:.4f}")
+
+test("mLSTM Block", test_mlstm_block)
+
+
+# ============================================================
+# Test 3: TMoE MLP
+# ============================================================
+def test_tmoe():
+    from vil_tracker.models.backbone import TMoEMLP
+    
+    tmoe = TMoEMLP(dim=384, mlp_ratio=4.0, num_experts=4)
+    params = count_params(tmoe)
+    print(f"  TMoEMLP params: {params:,} ({params/1e6:.3f}M)")
+    
+    x = torch.randn(2, 20, 384)
+    y = tmoe(x)
+    assert y.shape == (2, 20, 384), f"TMoE output shape: {y.shape}"
+    
+    # Test freezing shared expert
+    tmoe.freeze_shared_expert()
+    frozen = sum(1 for p in tmoe.shared_expert.parameters() if not p.requires_grad)
+    total_shared = sum(1 for p in tmoe.shared_expert.parameters())
+    assert frozen == total_shared, "Shared expert should be fully frozen"
+
+test("TMoE MLP", test_tmoe)
+
+
+# ============================================================
+# Test 4: Backbone (standard, small depth)
+# ============================================================
+def test_backbone_small():
+    from vil_tracker.models.backbone import ViLBackbone
+    
+    backbone = ViLBackbone(dim=384, depth=4, patch_size=16, tmoe_blocks=0)
+    params = count_params(backbone)
+    print(f"  Backbone (depth=4, no TMoE) params: {params:,} ({params/1e6:.3f}M)")
+    
+    template = torch.randn(2, 3, 128, 128)
+    search = torch.randn(2, 3, 256, 256)
+    
+    t_feat, s_feat = backbone(template, search)
+    assert t_feat.shape == (2, 64, 384), f"Template feat shape: {t_feat.shape}"
+    assert s_feat.shape == (2, 256, 384), f"Search feat shape: {s_feat.shape}"
+
+test("Backbone (standard, depth=4)", test_backbone_small)
+
+
+# ============================================================
+# Test 5: Backbone (with TMoE, depth=6)
+# ============================================================
+def test_backbone_tmoe():
+    from vil_tracker.models.backbone import ViLBackbone
+    
+    backbone = ViLBackbone(dim=384, depth=6, patch_size=16, tmoe_blocks=2, num_experts=4)
+    params = count_params(backbone)
+    print(f"  Backbone (depth=6, TMoE=2) params: {params:,} ({params/1e6:.3f}M)")
+    
+    template = torch.randn(1, 3, 128, 128)
+    search = torch.randn(1, 3, 256, 256)
+    
+    t_feat, s_feat = backbone(template, search)
+    assert t_feat.shape == (1, 64, 384), f"Template feat shape: {t_feat.shape}"
+    assert s_feat.shape == (1, 256, 384), f"Search feat shape: {s_feat.shape}"
+
+test("Backbone (with TMoE, depth=6)", test_backbone_tmoe)
+
+
+# ============================================================
+# Test 6: Prediction Heads
+# ============================================================
+def test_heads():
+    from vil_tracker.models.heads import CenterHead, UncertaintyHead, decode_predictions
+    
+    center_head = CenterHead(dim=384, feat_size=16)
+    unc_head = UncertaintyHead(dim=384, feat_size=16)
+    
+    print(f"  CenterHead params: {count_params(center_head):,}")
+    print(f"  UncertaintyHead params: {count_params(unc_head):,}")
+    
+    search_feat = torch.randn(2, 256, 384)
+    preds = center_head(search_feat)
+    
+    assert preds['heatmap'].shape == (2, 1, 16, 16), f"Heatmap shape: {preds['heatmap'].shape}"
+    assert preds['size'].shape == (2, 2, 16, 16), f"Size shape: {preds['size'].shape}"
+    assert preds['offset'].shape == (2, 2, 16, 16), f"Offset shape: {preds['offset'].shape}"
+    
+    # Decode
+    boxes, scores = decode_predictions(preds['heatmap'], preds['size'], preds['offset'])
+    assert boxes.shape == (2, 4), f"Boxes shape: {boxes.shape}"
+    assert scores.shape == (2,), f"Scores shape: {scores.shape}"
+    
+    # Uncertainty
+    log_var = unc_head(search_feat)
+    assert log_var.shape == (2, 1, 16, 16), f"Log variance shape: {log_var.shape}"
+
+test("Prediction Heads", test_heads)
+
+
+# ============================================================
+# Test 7: FiLM Temporal Modulation
+# ============================================================
+def test_film():
+    from vil_tracker.models.film_temporal import (
+        TemporalReliabilityCalibrator,
+        FiLMTemporalModulation,
+        TemporalModulationManager,
+    )
+    
+    # Test individual components
+    calib = TemporalReliabilityCalibrator(384)
+    film = FiLMTemporalModulation(384)
+    
+    x = torch.randn(2, 20, 384)
+    tc = torch.randn(2, 20, 384)
+    
+    rel = calib(tc)
+    assert rel.shape == (2, 20, 1), f"Reliability shape: {rel.shape}"
+    assert (rel >= 0).all() and (rel <= 1).all(), "Reliability not in [0,1]"
+    
+    modulated = film(x, tc, rel)
+    assert modulated.shape == (2, 20, 384), f"Modulated shape: {modulated.shape}"
+    
+    # Test manager
+    manager = TemporalModulationManager(dim=384, num_blocks=24, modulation_interval=6)
+    print(f"  TemporalModulationManager params: {count_params(manager):,}")
+    
+    # First call: no temporal context yet, should return unchanged
+    y = manager.modulate(x, block_idx=5)
+    assert torch.allclose(y, x), "Should return unchanged without temporal context"
+    
+    # Update context and try again
+    manager.update_temporal_context(x)
+    y = manager.modulate(x, block_idx=5)  # block 5 → (5+1)%6==0, should modulate
+    # With temporal context, output should differ
+    assert y.shape == (2, 20, 384)
+
+test("FiLM Temporal Modulation", test_film)
+
+
+# ============================================================
+# Test 8: Full Tracker (small depth for speed)
+# ============================================================
+def test_full_tracker_small():
+    from vil_tracker.models.tracker import ViLTracker, get_default_config
+    
+    config = get_default_config()
+    config['depth'] = 4
+    config['tmoe_blocks'] = 1
+    config['film_interval'] = 2
+    
+    tracker = ViLTracker(config)
+    params = count_params(tracker)
+    print(f"  Tracker (depth=4) params: {params:,} ({params/1e6:.3f}M)")
+    
+    template = torch.randn(2, 3, 128, 128)
+    search = torch.randn(2, 3, 256, 256)
+    
+    output = tracker(template, search)
+    
+    assert output['heatmap'].shape == (2, 1, 16, 16)
+    assert output['size'].shape == (2, 2, 16, 16)
+    assert output['boxes'].shape == (2, 4)
+    assert output['scores'].shape == (2,)
+    assert 'log_variance' in output
+    
+    print(f"  Predicted boxes: {output['boxes'][0].tolist()}")
+    print(f"  Scores: {output['scores'].tolist()}")
+
+test("Full Tracker (depth=4)", test_full_tracker_small)
+
+
+# ============================================================
+# Test 9: Loss Functions
+# ============================================================
+def test_losses():
+    from vil_tracker.training.losses import (
+        FocalLoss, GIoULoss, UncertaintyNLLLoss,
+        MemoryContrastiveLoss, CombinedTrackingLoss,
+    )
+    
+    B = 4
+    
+    # Focal loss
+    focal = FocalLoss()
+    pred_hm = torch.randn(B, 1, 16, 16)
+    gt_hm = torch.zeros(B, 1, 16, 16)
+    gt_hm[:, :, 8, 8] = 1.0
+    fl = focal(pred_hm, gt_hm)
+    print(f"  Focal loss: {fl.item():.4f}")
+    assert fl.item() > 0, "Focal loss should be positive"
+    
+    # GIoU loss
+    giou = GIoULoss()
+    pred_box = torch.tensor([[128.0, 128.0, 50.0, 50.0]] * B)
+    gt_box = torch.tensor([[130.0, 130.0, 48.0, 48.0]] * B)
+    gl = giou(pred_box, gt_box)
+    print(f"  GIoU loss: {gl.item():.4f}")
+    assert 0 <= gl.item() <= 2, f"GIoU loss out of range: {gl.item()}"
+    
+    # Contrastive loss
+    contrastive = MemoryContrastiveLoss()
+    feat_a = torch.randn(B, 384)
+    feat_b = feat_a + torch.randn(B, 384) * 0.1  # slightly perturbed
+    cl = contrastive(feat_a, feat_b)
+    print(f"  Contrastive loss: {cl.item():.4f}")
+    
+    # Combined loss
+    combined = CombinedTrackingLoss()
+    pred = {
+        'heatmap': pred_hm,
+        'size': torch.rand(B, 2, 16, 16),
+        'boxes': pred_box,
+        'log_variance': torch.randn(B, 1, 16, 16),
+    }
+    loss_dict = combined(pred, gt_hm, torch.tensor([[0.2, 0.2]] * B), gt_box)
+    print(f"  Combined loss: {loss_dict['total'].item():.4f}")
+    assert loss_dict['total'].item() > 0
+
+test("Loss Functions", test_losses)
+
+
+# ============================================================
+# Test 10: Kalman Filter
+# ============================================================
+def test_kalman():
+    from vil_tracker.inference.kalman import KalmanFilter
+    
+    kf = KalmanFilter()
+    assert not kf.initialized
+    
+    # Initialize
+    init_box = np.array([100.0, 100.0, 50.0, 50.0])
+    kf.initialize(init_box)
+    assert kf.initialized
+    
+    # Predict + update cycle
+    for i in range(10):
+        pred = kf.predict()
+        assert len(pred) == 4, f"Prediction length: {len(pred)}"
+        
+        # Simulate noisy measurement
+        noise = np.random.randn(4) * 2
+        meas = init_box + np.array([i * 2, i * 1, 0, 0]) + noise
+        kf.update(meas, uncertainty=1.0)
+    
+    state = kf.get_state()
+    print(f"  Final state: cx={state[0]:.1f}, cy={state[1]:.1f}, w={state[2]:.1f}, h={state[3]:.1f}")
+    assert state[2] > 0 and state[3] > 0, "Width/height should be positive"
+
+test("Kalman Filter", test_kalman)
+
+
+# ============================================================
+# Test 11: Dataset (synthetic)
+# ============================================================
+def test_dataset():
+    from vil_tracker.data.dataset import TrackingDataset
+    
+    ds = TrackingDataset(synthetic=True, synthetic_length=100)
+    assert len(ds) == 100
+    
+    sample = ds[0]
+    assert sample['template'].shape == (3, 128, 128), f"Template shape: {sample['template'].shape}"
+    assert sample['search'].shape == (3, 256, 256), f"Search shape: {sample['search'].shape}"
+    assert sample['heatmap'].shape == (1, 16, 16), f"Heatmap shape: {sample['heatmap'].shape}"
+    assert sample['size'].shape == (2,), f"Size shape: {sample['size'].shape}"
+    assert sample['boxes'].shape == (4,), f"Boxes shape: {sample['boxes'].shape}"
+    
+    # Check ACL difficulty changes output
+    ds.set_acl_difficulty(0.0)
+    easy_sample = ds[42]
+    ds.set_acl_difficulty(1.0)
+    hard_sample = ds[42]
+    print(f"  Easy center: {easy_sample['boxes'][:2].tolist()}")
+    print(f"  Hard center: {hard_sample['boxes'][:2].tolist()}")
+
+test("Dataset (synthetic)", test_dataset)
+
+
+# ============================================================
+# Test 12: Training Step (mini forward + backward)
+# ============================================================
+def test_training_step():
+    from vil_tracker.models.tracker import ViLTracker, get_default_config
+    from vil_tracker.training.losses import CombinedTrackingLoss
+    from vil_tracker.models.heads import generate_heatmap
+    
+    config = get_default_config()
+    config['depth'] = 2
+    config['tmoe_blocks'] = 0
+    config['film_interval'] = 2
+    
+    model = ViLTracker(config)
+    model.train()
+    loss_fn = CombinedTrackingLoss()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
+    
+    B = 2
+    template = torch.randn(B, 3, 128, 128)
+    search = torch.randn(B, 3, 256, 256)
+    
+    # GT targets
+    gt_center = torch.tensor([[128.0, 128.0], [100.0, 150.0]])
+    gt_heatmap = generate_heatmap(gt_center, feat_size=16, search_size=256)
+    gt_size = torch.tensor([[0.2, 0.3], [0.15, 0.25]])
+    gt_boxes = torch.tensor([[128.0, 128.0, 51.2, 76.8], [100.0, 150.0, 38.4, 64.0]])
+    
+    # Forward
+    pred = model(template, search)
+    loss_dict = loss_fn(pred, gt_heatmap, gt_size, gt_boxes)
+    
+    # Backward
+    loss_dict['total'].backward()
+    
+    # Check gradients exist
+    has_grads = sum(1 for p in model.parameters() if p.grad is not None)
+    total_params_count = sum(1 for p in model.parameters())
+    print(f"  Loss: {loss_dict['total'].item():.4f}")
+    print(f"  Params with gradients: {has_grads}/{total_params_count}")
+    
+    # Optimizer step
+    optimizer.step()
+    optimizer.zero_grad()
+    
+    assert loss_dict['total'].item() > 0
+    assert has_grads > 0
+
+test("Training Step (depth=2)", test_training_step)
+
+
+# ============================================================
+# Test 13: Model Summary (FULL depth=24, constraint check)
+# ============================================================
+def test_model_summary():
+    from vil_tracker.models.tracker import ViLTracker, get_default_config
+    from vil_tracker.utils.helpers import print_model_summary
+    
+    config = get_default_config()
+    model = ViLTracker(config)
+    
+    summary = print_model_summary(model, config)
+    
+    total_m = summary['total_params'] / 1e6
+    
+    # HARD CONSTRAINTS
+    assert summary['param_ok'], f"FAIL: {total_m:.2f}M params exceeds 50M limit"
+    assert summary['size_ok'], f"FAIL: {summary['size_fp16_mb']:.1f}MB exceeds 500MB limit"
+    # GFLOPs is approximate, warn but don't fail if close
+    if not summary['flop_ok']:
+        print(f"  ⚠️  GFLOPs estimate ({summary['gflops']:.2f}) exceeds 20, but this is approximate")
+
+test("Model Summary (full depth=24)", test_model_summary)
+
+
+# ============================================================
+# Summary
+# ============================================================
+print("\n" + "=" * 60)
+print(f"Results: {PASS}/{PASS + FAIL} tests passed")
+if FAIL > 0:
+    print(f"  ❌ {FAIL} test(s) FAILED")
+    sys.exit(1)
+else:
+    print("  ✅ All tests passed!")
+    sys.exit(0)