Coverage for cafl4ds/models/heads.py: 96%

1"""Task heads attached to the shared :class:`~cafl4ds.models.vit.TinyViTEncoder`. 

2 

3These are *training* heads, discarded at measurement time (the health instruments read the 

4encoder's pooled representation, never these outputs): 

5 

6* :class:`MLPHead` — a configurable BN-MLP; used as SimSiam's projector and predictor 

7 (BYOL later reuses the same shape). 

8* :class:`MAEDecoder` — a lightweight transformer decoder that reconstructs pixels from the 

9 encoder's visible-token latent (He et al. 2022). 

10""" 

11 

12from __future__ import annotations 

13 

14from typing import cast 

15 

16import torch 

17from torch import nn 

18 

19from cafl4ds.models.vit import Block 

20 

21 

22class MLPHead(nn.Module): # type: ignore[misc] # nn.Module is Any without torch stubs (mypy hook env) 

23 """A batch-norm MLP head (SimSiam/BYOL projector or predictor). 

24 

25 With ``num_layers=3`` and ``last_bn=True`` this is the SimSiam projector; with 

26 ``num_layers=2`` and ``last_bn=False`` it is the SimSiam predictor (Chen & He 2021). 

27 """ 

28 

29 def __init__( 

30 self, 

31 in_dim: int, 

32 hidden_dim: int, 

33 out_dim: int, 

34 num_layers: int = 3, 

35 last_bn: bool = True, 

36 ) -> None: 

37 """Build the MLP head. 

38 

39 Args: 

40 in_dim: Input dimensionality. 

41 hidden_dim: Hidden-layer width (used for all but the output layer). 

42 out_dim: Output dimensionality. 

43 num_layers: Total number of linear layers (``>= 1``). 

44 last_bn: Whether to apply a (bias-free) BatchNorm after the output layer, with 

45 no final activation — the SimSiam projector does, the predictor does not. 

46 

47 Raises: 

48 ValueError: If ``num_layers < 1``. 

49 """ 

50 super().__init__() 

51 if num_layers < 1: 51 ↛ 52line 51 didn't jump to line 52 because the condition on line 51 was never true

52 raise ValueError(f"num_layers must be >= 1; got {num_layers}.") 

53 layers: list[nn.Module] = [] 

54 dim = in_dim 

55 for _ in range(num_layers - 1): 

56 layers += [nn.Linear(dim, hidden_dim, bias=False), nn.BatchNorm1d(hidden_dim), nn.ReLU(inplace=True)] 

57 dim = hidden_dim 

58 layers.append(nn.Linear(dim, out_dim, bias=not last_bn)) 

59 if last_bn: 

60 layers.append(nn.BatchNorm1d(out_dim, affine=False)) 

61 self.net = nn.Sequential(*layers) 

62 

63 def forward(self, x: torch.Tensor) -> torch.Tensor: 

64 """Apply the MLP head (``[B, in_dim] -> [B, out_dim]``).""" 

65 return cast(torch.Tensor, self.net(x)) 

66 

67 

68class MAEDecoder(nn.Module): # type: ignore[misc] # nn.Module is Any without torch stubs (mypy hook env) 

69 """A lightweight MAE decoder: reconstruct pixels from the visible-token latent. 

70 

71 Following He et al. (2022): project the encoder latent into the decoder width, scatter 

72 the visible tokens back into the full sequence (filling the holes with a shared learned 

73 mask token), add decoder positional embeddings, run a few transformer blocks, and 

74 linearly predict the per-patch pixels. 

75 """ 

76 

77 def __init__( 

78 self, 

79 num_patches: int, 

80 encoder_dim: int, 

81 patch_size: int, 

82 in_chans: int = 3, 

83 decoder_dim: int = 64, 

84 depth: int = 2, 

85 num_heads: int = 3, 

86 mlp_ratio: float = 2.0, 

87 ) -> None: 

88 """Build the decoder. 

89 

90 Args: 

91 num_patches: Number of patches ``N`` in the full (unmasked) sequence. 

92 encoder_dim: Dimensionality of the encoder latent tokens. 

93 patch_size: Side length ``p`` of each square patch (sets the pixel output width). 

94 in_chans: Number of image channels ``C``. 

95 decoder_dim: Decoder token dimensionality (smaller than the encoder's). 

96 depth: Number of decoder transformer blocks. 

97 num_heads: Number of attention heads per decoder block. 

98 mlp_ratio: Decoder MLP hidden width as a multiple of ``decoder_dim``. 

99 """ 

100 super().__init__() 

101 self.decoder_embed = nn.Linear(encoder_dim, decoder_dim) 

102 self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim)) 

103 self.decoder_pos_embed = nn.Parameter(torch.zeros(1, 1 + num_patches, decoder_dim)) 

104 self.blocks = nn.ModuleList([Block(decoder_dim, num_heads, mlp_ratio) for _ in range(depth)]) 

105 self.norm = nn.LayerNorm(decoder_dim) 

106 self.pred = nn.Linear(decoder_dim, patch_size * patch_size * in_chans) 

107 nn.init.trunc_normal_(self.mask_token, std=0.02) 

108 nn.init.trunc_normal_(self.decoder_pos_embed, std=0.02) 

109 

110 def forward(self, latent: torch.Tensor, ids_restore: torch.Tensor) -> torch.Tensor: 

111 """Reconstruct per-patch pixels from the encoder's visible-token latent. 

112 

113 Args: 

114 latent: Encoder output ``[B, 1 + N_keep, encoder_dim]`` (``cls`` token first). 

115 ids_restore: The ``[B, N]`` permutation from 

116 :meth:`~cafl4ds.models.vit.TinyViTEncoder.random_masking` that undoes the 

117 visible-patch shuffle. 

118 

119 Returns: 

120 Predicted patches ``[B, N, p * p * C]`` in row-major patch order (``cls`` token 

121 dropped). 

122 """ 

123 x = self.decoder_embed(latent) 

124 b, n = ids_restore.shape 

125 n_keep = x.shape[1] - 1 # exclude cls 

126 mask_tokens = self.mask_token.expand(b, n - n_keep, -1) 

127 x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1) # drop cls, append mask tokens 

128 x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).expand(-1, -1, x.shape[2])) # unshuffle 

129 x = torch.cat([x[:, :1, :], x_], dim=1) # restore cls 

130 x = x + self.decoder_pos_embed 

131 for blk in self.blocks: 

132 x = blk(x) 

133 x = self.norm(x) 

134 return cast(torch.Tensor, self.pred(x[:, 1:, :])) # drop cls; [B, N, p*p*C]