(test_data, sh_degree: int)
| 579 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") |
| 580 | @pytest.mark.parametrize("sh_degree", [0, 1, 2, 3, 4]) |
| 581 | def test_sh(test_data, sh_degree: int): |
| 582 | from gsplat.cuda._torch_impl import _spherical_harmonics |
| 583 | from gsplat.cuda._wrapper import spherical_harmonics |
| 584 | |
| 585 | torch.manual_seed(42) |
| 586 | |
| 587 | N = 1000 |
| 588 | coeffs = torch.randn(N, (4 + 1) ** 2, 3, device=device) |
| 589 | dirs = torch.randn(N, 3, device=device) |
| 590 | coeffs.requires_grad = True |
| 591 | dirs.requires_grad = True |
| 592 | |
| 593 | colors = spherical_harmonics(sh_degree, dirs, coeffs) |
| 594 | _colors = _spherical_harmonics(sh_degree, dirs, coeffs) |
| 595 | torch.testing.assert_close(colors, _colors, rtol=1e-4, atol=1e-4) |
| 596 | |
| 597 | v_colors = torch.randn_like(colors) |
| 598 | |
| 599 | v_coeffs, v_dirs = torch.autograd.grad( |
| 600 | (colors * v_colors).sum(), (coeffs, dirs), retain_graph=True, allow_unused=True |
| 601 | ) |
| 602 | _v_coeffs, _v_dirs = torch.autograd.grad( |
| 603 | (_colors * v_colors).sum(), (coeffs, dirs), retain_graph=True, allow_unused=True |
| 604 | ) |
| 605 | torch.testing.assert_close(v_coeffs, _v_coeffs, rtol=1e-4, atol=1e-4) |
| 606 | if sh_degree > 0: |
| 607 | torch.testing.assert_close(v_dirs, _v_dirs, rtol=1e-4, atol=1e-4) |
nothing calls this directly
no test coverage detected