| """Real Metal tests; teacher uses random tiny modules, never pretrained weights.""" |
| import pytest |
| import torch |
|
|
| from image3d_student.model import Student |
| from image3d_student.teacher_idol import patch_attention |
|
|
| pytestmark = pytest.mark.skipif(not torch.backends.mps.is_available(), reason="Metal GPU not visible") |
|
|
|
|
| def test_default_student_backward_on_metal(): |
| model = Student().to("mps") |
| result = model(torch.rand(1,3,128,128,device="mps")) |
| loss = sum(t.square().mean() for t in result.values()) |
| loss.backward() |
| assert torch.isfinite(loss) |
| assert result["xyz"].shape == (1,4096,3) |
| assert all(p.grad is None or torch.isfinite(p.grad).all() for p in model.parameters()) |
|
|
|
|
| @torch.no_grad() |
| def test_tiny_real_idol_architecture_on_metal(): |
| pytest.importorskip("idol") |
| from tensordict import TensorDict |
| from idol.models import LatentSeqMMFlowModel, OctreeProbabilityFixedlenDecoder, GaussianFixedlenDecoder, layers |
| from idol.sampler import sample_flow |
| patch_attention(layers, query_chunk=4) |
| model = LatentSeqMMFlowModel(q_token_length=8,in_channels=4,cam_channels=6, |
| model_channels=32,cond_channels=12,cond2_channels=8,num_blocks=1, |
| num_refiner_blocks=2,num_heads=2,num_head_channels=16).eval().to("mps") |
| noise = TensorDict({"latent": torch.randn(1,8,4,device="mps"), |
| "camera": torch.randn(1,1,6,device="mps")}, batch_size=1, device="mps") |
| cond = TensorDict({"feature1": torch.randn(1,6,12,device="mps"), |
| "feature2": torch.randn(1,6,8,device="mps")}, batch_size=1, device="mps") |
| result = sample_flow(model,noise,cond,cond.apply(torch.zeros_like),steps=2,verbose=False,batch_cfg=False) |
| latent = result["latent"] |
| assert torch.isfinite(latent).all() |
| octree = OctreeProbabilityFixedlenDecoder(32,4,1,2).eval().to("mps") |
| points = octree.sample(octree,latent,num_points=8,level=3) |
| assert points.shape == (1,8,3) |
| config = {"num_gaussians": 4, "lr": {k: 1. for k in ["_xyz","_features_dc","_scaling","_rotation","_opacity"]}, |
| "perturbe_size": .1, "scaling_bias": .01, "opacity_bias": .8, "3d_filter_kernel_size": .001} |
| decoder = GaussianFixedlenDecoder(3,32,4,1,2,config).eval().to("mps") |
| gaussian = decoder.to_representation(points,decoder(points,latent))[0] |
| assert gaussian.xyz.shape == (32,3) |
| for value in (gaussian.xyz,gaussian.features_dc,gaussian.scaling,gaussian.rotation,gaussian.opacity): |
| assert torch.isfinite(value).all() |
|
|
