| """ |
| Sanity check: verify teacher and student forward passes work correctly. |
| Run on CPU to test before GPU training. |
| """ |
| import torch |
| from models.teacher import PriviGazeTeacher |
| from models.student import PriviGazeStudent, count_parameters |
| from models.distillation_loss import PriviGazeDistillationLoss |
| from models.dataset import SyntheticGazeDataset |
|
|
|
|
| def test_teacher(): |
| print("Testing Teacher...") |
| model = PriviGazeTeacher() |
| le = torch.randn(2, 3, 112, 112) |
| re = torch.randn(2, 3, 112, 112) |
| fb = torch.randn(2, 1, 224, 224) |
| pp, yp, pl, yl, feat = model(le, re, fb) |
| assert pp.shape == (2,), f"pitch shape {pp.shape} != (2,)" |
| assert yp.shape == (2,), f"yaw shape {yp.shape} != (2,)" |
| assert pl.shape == (2, 90), f"pitch_logits {pl.shape} != (2,90)" |
| assert yl.shape == (2, 90), f"yaw_logits {yl.shape} != (2,90)" |
| assert feat.shape == (2, 256), f"features {feat.shape} != (2,256)" |
| print(f" OK! Params: {count_parameters(model):,}") |
|
|
|
|
| def test_student(): |
| print("Testing Student...") |
| model = PriviGazeStudent() |
| fg = torch.randn(2, 1, 224, 224) |
| pp, yp, feat = model(fg) |
| assert pp.shape == (2,), f"pitch shape {pp.shape}" |
| assert yp.shape == (2,), f"yaw shape {yp.shape}" |
| assert feat.shape == (2, 128), f"features {feat.shape}" |
| print(f" OK! Params: {count_parameters(model):,}") |
|
|
|
|
| def test_distillation_loss(): |
| print("Testing Distillation Loss...") |
| loss_fn = PriviGazeDistillationLoss() |
| sp = torch.randn(4) |
| sy = torch.randn(4) |
| spl = torch.randn(4, 90) |
| syl = torch.randn(4, 90) |
| sf = torch.randn(4, 128) |
| tp = torch.randn(4) |
| ty = torch.randn(4) |
| tpl = torch.randn(4, 90) |
| tyl = torch.randn(4, 90) |
| tf = torch.randn(4, 256) |
| pt = torch.randn(4) |
| yt = torch.randn(4) |
| loss, ld = loss_fn(sp, sy, spl, syl, sf, tp, ty, tpl, tyl, tf, pt, yt) |
| assert loss.item() > 0, "loss should be positive" |
| assert 'loss_total' in ld |
| print(f" OK! Loss: {loss.item():.4f}") |
|
|
|
|
| def test_dataset(): |
| print("Testing Dataset...") |
| ds = SyntheticGazeDataset(num_samples=10) |
| sample = ds[0] |
| assert sample['left_eye'].shape == (3, 112, 112) |
| assert sample['right_eye'].shape == (3, 112, 112) |
| assert sample['face_blurred_gray'].shape == (1, 224, 224) |
| assert sample['face_gray'].shape == (1, 224, 224) |
| assert -90 <= sample['pitch'].item() <= 90 |
| assert -90 <= sample['yaw'].item() <= 90 |
| print(f" OK! {len(ds)} samples") |
|
|
|
|
| def test_end_to_end(): |
| print("Testing End-to-End (1 batch)...") |
| teacher = PriviGazeTeacher() |
| student = PriviGazeStudent() |
| loss_fn = PriviGazeDistillationLoss() |
| ds = SyntheticGazeDataset(num_samples=4) |
| batch = {k: torch.stack([ds[i][k] for i in range(4)]) for k in ds[0].keys()} |
|
|
| le = batch['left_eye'] |
| re = batch['right_eye'] |
| fb = batch['face_blurred_gray'] |
| fg = batch['face_gray'] |
| pt = batch['pitch'] |
| yt = batch['yaw'] |
|
|
| with torch.no_grad(): |
| tp, ty, tpl, tyl, tf = teacher(le, re, fb) |
| sp, sy, sf = student(fg) |
| spl = student.pitch_head(sf) |
| syl = student.yaw_head(sf) |
|
|
| loss, ld = loss_fn(sp, sy, spl, syl, sf, tp, ty, tpl, tyl, tf, pt, yt) |
| loss.backward() |
| print(f" OK! Backward pass succeeded. Loss: {loss.item():.4f}") |
|
|
|
|
| if __name__ == "__main__": |
| test_teacher() |
| test_student() |
| test_distillation_loss() |
| test_dataset() |
| test_end_to_end() |
| print("\nAll tests passed!") |
|
|
