Coverage for cafl4ds/ssl/mae.py: 95%

1"""Masked Autoencoder (MAE) SSL method — the default ``C`` backbone (He et al. 2022). 

2 

3MAE masks a large fraction of image patches, encodes the visible ones, and reconstructs the 

4missing pixels; the per-patch reconstruction error is a free label-free informativeness 

5signal (used by novelty filters in later phases). MAE is **collapse-resistant** — a constant 

6code cannot reconstruct diverse pixels — so its Phase-1 degradation mode is 

7forgetting/overspecialization, not representational collapse; the collapse demonstration uses 

8the joint-embedding method (:mod:`cafl4ds.ssl.simsiam`) instead. 

9""" 

10 

11from __future__ import annotations 

12 

13from typing import cast 

14 

15import torch 

16from torchvision.transforms import v2 

17 

18from cafl4ds.models.heads import MAEDecoder 

19from cafl4ds.models.vit import TinyViTEncoder, patchify 

20from cafl4ds.ssl.augment import make_light_augment 

21from cafl4ds.ssl.base import SSLMethod 

22 

23 

24class MAE(SSLMethod): 

25 """Masked Autoencoder over the shared :class:`~cafl4ds.models.vit.TinyViTEncoder`.""" 

26 

27 def __init__( 

28 self, 

29 encoder: TinyViTEncoder, 

30 decoder: MAEDecoder, 

31 mask_ratio: float = 0.75, 

32 norm_pix_loss: bool = True, 

33 augment: v2.Transform | None = None, 

34 ) -> None: 

35 """Build the MAE method. 

36 

37 Args: 

38 encoder: The shared backbone encoder. 

39 decoder: The pixel-reconstruction decoder. 

40 mask_ratio: Fraction of patches masked each step. 

41 norm_pix_loss: Whether to normalize each target patch (per-patch mean/var) before 

42 the reconstruction loss, as in the MAE paper (usually more stable). 

43 augment: Light spatial augmentation applied before masking; defaults to 

44 :func:`~cafl4ds.ssl.augment.make_light_augment` sized to the encoder. 

45 """ 

46 super().__init__(encoder) 

47 self.decoder = decoder 

48 self.mask_ratio = mask_ratio 

49 self.norm_pix_loss = norm_pix_loss 

50 img_size = int(round(encoder.num_patches**0.5)) * encoder.patch_size 

51 self.augment = augment if augment is not None else make_light_augment(img_size) 

52 

53 @property 

54 def name(self) -> str: 

55 """Return the method identifier.""" 

56 return "mae" 

57 

58 def _reconstruction_target(self, imgs: torch.Tensor) -> torch.Tensor: 

59 """Patchify images into reconstruction targets, optionally per-patch normalized. 

60 

61 Args: 

62 imgs: Images ``[B, C, H, W]``. 

63 

64 Returns: 

65 Target patches ``[B, N, p * p * C]``. 

66 """ 

67 target = patchify(imgs, self.encoder.patch_size) 

68 if self.norm_pix_loss: 68 ↛ 72line 68 didn't jump to line 72 because the condition on line 68 was always true

69 mean = target.mean(dim=-1, keepdim=True) 

70 var = target.var(dim=-1, keepdim=True) 

71 target = (target - mean) / (var + 1.0e-6).sqrt() 

72 return target 

73 

74 def _masked_recon_loss(self, imgs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: 

75 """Augment, mask, reconstruct, and return the per-patch loss and the mask. 

76 

77 Shared by :meth:`training_step` (mean over masked patches) and 

78 :meth:`per_sample_loss` (per-image mean over its masked patches). Masking is random, 

79 so each call draws a fresh mask — fine for both a training step and a selection probe. 

80 

81 Args: 

82 imgs: A batch of raw images ``[B, C, H, W]``. 

83 

84 Returns: 

85 A ``(loss_per_patch, mask)`` pair, each ``[B, N]``; ``mask`` is 1 on masked 

86 (reconstructed) patches and 0 on visible ones. 

87 """ 

88 views = self.augment(imgs) 

89 latent, mask, ids_restore = self.encoder.forward_encoder(views, mask_ratio=self.mask_ratio) 

90 assert mask is not None and ids_restore is not None # noqa: S101 - guaranteed by mask_ratio > 0 

91 pred = self.decoder(latent, ids_restore) 

92 target = self._reconstruction_target(views) 

93 loss_per_patch = (pred - target).pow(2).mean(dim=-1) # [B, N] 

94 return loss_per_patch, mask 

95 

96 def training_step(self, imgs: torch.Tensor) -> torch.Tensor: 

97 """Compute the masked reconstruction loss on the masked patches only. 

98 

99 Args: 

100 imgs: A batch of raw images ``[B, C, H, W]``. 

101 

102 Returns: 

103 The mean squared reconstruction error over masked patches (scalar). 

104 """ 

105 loss_per_patch, mask = self._masked_recon_loss(imgs) 

106 return cast(torch.Tensor, (loss_per_patch * mask).sum() / mask.sum().clamp_min(1.0)) 

107 

108 def per_sample_loss(self, imgs: torch.Tensor) -> torch.Tensor: 

109 """Per-image masked reconstruction MSE ``[B]`` (no gradient). 

110 

111 The free per-frame informativeness signal MAE exposes: high loss = the current model 

112 reconstructs this frame poorly. Averaged over each image's own masked patches. 

113 

114 Args: 

115 imgs: A batch of raw images ``[B, C, H, W]``. 

116 

117 Returns: 

118 A detached ``[B]`` tensor of per-image reconstruction errors. 

119 """ 

120 with torch.no_grad(): 

121 loss_per_patch, mask = self._masked_recon_loss(imgs) 

122 per_image = (loss_per_patch * mask).sum(dim=-1) / mask.sum(dim=-1).clamp_min(1.0) 

123 return per_image