| |
|
|
| |
|
|
| """Tests for the PyTorch reference SAM3 pipeline.""" |
|
|
| import pytest |
| import torch |
|
|
|
|
| class TestSam3ReferenceModel: |
| """Verify SAM3 PyTorch model loads and runs correctly.""" |
|
|
| def test_model_loads(self, sam3_reference_model): |
| """SAM3 model should load without errors.""" |
| assert sam3_reference_model is not None |
| assert not sam3_reference_model.training |
|
|
| def test_vit_backbone_structure(self, sam3_vit_backbone): |
| """Verify ViT backbone has expected structure.""" |
| vit = sam3_vit_backbone |
|
|
| |
| assert hasattr(vit, "patch_embed") |
| assert vit.patch_embed.proj.weight.shape == (1024, 3, 14, 14) |
|
|
| |
| assert len(vit.blocks) == 32 |
|
|
| |
| block = vit.blocks[0] |
| assert hasattr(block, "norm1") |
| assert hasattr(block, "attn") |
| assert hasattr(block, "norm2") |
| assert hasattr(block, "mlp") |
|
|
| |
| assert block.attn.qkv.weight.shape == (3072, 1024) |
| assert block.attn.num_heads == 16 |
|
|
| |
| assert block.mlp.fc1.weight.shape[0] == 4736 |
| assert block.mlp.fc2.weight.shape[1] == 4736 |
|
|
| def test_vit_backbone_window_vs_global(self, sam3_vit_backbone): |
| """Verify correct blocks use window vs global attention.""" |
| global_blocks = {7, 15, 23, 31} |
| for i, block in enumerate(sam3_vit_backbone.blocks): |
| if i in global_blocks: |
| assert block.window_size == 0, f"Block {i} should use global attention" |
| else: |
| assert block.window_size == 24, f"Block {i} should use window attention (ws=24)" |
|
|
| def test_vit_backbone_forward_shapes(self, sam3_vit_backbone): |
| """Test ViT backbone produces correct output shapes.""" |
| torch.manual_seed(42) |
| pixel_values = torch.randn(1, 3, 1008, 1008) |
|
|
| with torch.no_grad(): |
| outputs = sam3_vit_backbone(pixel_values) |
|
|
| |
| |
| assert isinstance(outputs, list), f"Expected list, got {type(outputs)}" |
| assert len(outputs) >= 1, f"Expected at least 1 output, got {len(outputs)}" |
|
|
| |
| feats = outputs[-1] |
| assert feats.shape[0] == 1, f"Batch dim mismatch: {feats.shape}" |
| assert feats.shape[1] == 1024, f"Channel dim should be 1024, got {feats.shape[1]}" |
| assert feats.shape[2] == 72, f"Height should be 72, got {feats.shape[2]}" |
| assert feats.shape[3] == 72, f"Width should be 72, got {feats.shape[3]}" |
|
|
| def test_patch_embed_output(self, sam3_vit_backbone): |
| """Test PatchEmbed produces correct shape.""" |
| torch.manual_seed(42) |
| pixel_values = torch.randn(1, 3, 1008, 1008) |
|
|
| with torch.no_grad(): |
| patch_output = sam3_vit_backbone.patch_embed(pixel_values) |
|
|
| |
| assert patch_output.shape == (1, 72, 72, 1024), f"Unexpected shape: {patch_output.shape}" |
|
|
| def test_single_block_forward(self, sam3_vit_backbone): |
| """Test a single ViT block produces correct output shape.""" |
| torch.manual_seed(42) |
| block = sam3_vit_backbone.blocks[0] |
|
|
| |
| hidden_states = torch.randn(1, 72, 72, 1024) |
|
|
| with torch.no_grad(): |
| output = block(hidden_states) |
|
|
| assert output.shape == (1, 72, 72, 1024), f"Unexpected shape: {output.shape}" |
|
|
| def test_vit_rope_present(self, sam3_vit_backbone): |
| """Verify RoPE is configured in attention blocks.""" |
| block = sam3_vit_backbone.blocks[0] |
| assert block.attn.use_rope, "Block should use RoPE" |
| assert block.attn.freqs_cis is not None, "RoPE frequencies should be precomputed" |
|
|
| def test_extract_block_params(self, sam3_vit_backbone): |
| """Test parameter extraction from ViT block.""" |
| from tt_sam3.reference.torch_sam3 import extract_vit_block_params |
|
|
| params = extract_vit_block_params(sam3_vit_backbone.blocks[0]) |
|
|
| assert "norm1_weight" in params |
| assert "attn_qkv_weight" in params |
| assert "mlp_fc1_weight" in params |
| assert params["attn_qkv_weight"].shape == (3072, 1024) |
| assert params["mlp_fc1_weight"].shape == (4736, 1024) |
| assert params["window_size"] == 24 |
|
|
| def test_neck_structure(self, sam3_neck): |
| """Verify FPN neck has expected structure.""" |
| assert len(sam3_neck.convs) == 4 |
|
|
| def test_transformer_structure(self, sam3_transformer): |
| """Verify transformer encoder/decoder structure.""" |
| assert sam3_transformer.d_model == 256 |
| assert hasattr(sam3_transformer, "encoder") |
| assert hasattr(sam3_transformer, "decoder") |
|
|
