Coverage for cafl4ds/ssl/factory.py: 100%

1"""Factories that build an :class:`~cafl4ds.ssl.base.SSLMethod` from an encoder + scalars. 

2 

3The SSL methods take *pre-built* heads (decoder / projector / predictor) so they stay 

4decoupled from any particular sizing. These factories bridge the gap for Hydra: a config 

5supplies plain scalar hyper-parameters, the run script supplies the shared ``encoder``, and 

6the factory sizes the heads from the encoder's own dimensions. This keeps the ``ssl`` config 

7group a flat list of numbers while the head shapes stay consistent with the backbone. 

8""" 

9 

10from __future__ import annotations 

11 

12from cafl4ds.models.heads import MAEDecoder, MLPHead 

13from cafl4ds.models.vit import TinyViTEncoder 

14from cafl4ds.ssl.mae import MAE 

15from cafl4ds.ssl.simsiam import SimSiam 

16 

17 

18def build_mae( 

19 encoder: TinyViTEncoder, 

20 mask_ratio: float = 0.75, 

21 decoder_dim: int = 64, 

22 decoder_depth: int = 2, 

23 decoder_heads: int = 4, 

24 decoder_mlp_ratio: float = 2.0, 

25 norm_pix_loss: bool = True, 

26) -> MAE: 

27 """Build an :class:`~cafl4ds.ssl.mae.MAE` with a decoder sized to the encoder. 

28 

29 Args: 

30 encoder: The shared backbone encoder. 

31 mask_ratio: Fraction of patches masked each step. 

32 decoder_dim: Decoder token width. 

33 decoder_depth: Number of decoder transformer blocks. 

34 decoder_heads: Attention heads per decoder block (must divide ``decoder_dim``). 

35 decoder_mlp_ratio: Decoder MLP hidden width as a multiple of ``decoder_dim``. 

36 norm_pix_loss: Whether to per-patch normalize the reconstruction targets. 

37 

38 Returns: 

39 The assembled MAE method. 

40 """ 

41 decoder = MAEDecoder( 

42 num_patches=encoder.num_patches, 

43 encoder_dim=encoder.embed_dim, 

44 patch_size=encoder.patch_size, 

45 in_chans=encoder.in_chans, 

46 decoder_dim=decoder_dim, 

47 depth=decoder_depth, 

48 num_heads=decoder_heads, 

49 mlp_ratio=decoder_mlp_ratio, 

50 ) 

51 return MAE(encoder, decoder, mask_ratio=mask_ratio, norm_pix_loss=norm_pix_loss) 

52 

53 

54def build_simsiam( 

55 encoder: TinyViTEncoder, 

56 proj_hidden: int = 256, 

57 proj_dim: int = 128, 

58 pred_hidden: int = 64, 

59 anti_collapse: bool = True, 

60) -> SimSiam: 

61 """Build a :class:`~cafl4ds.ssl.simsiam.SimSiam` with heads sized to the encoder. 

62 

63 Args: 

64 encoder: The shared backbone encoder. 

65 proj_hidden: Hidden width of the 3-layer projector. 

66 proj_dim: Output width of the projector (and the predictor's input/output). 

67 pred_hidden: Hidden width of the 2-layer predictor bottleneck. 

68 anti_collapse: Keep SimSiam's predictor + stop-gradient (``True``, the healthy 

69 control) or disable both for the forced-collapse positive control (``False``). 

70 

71 Returns: 

72 The assembled SimSiam method. 

73 """ 

74 projector = MLPHead(encoder.embed_dim, proj_hidden, proj_dim, num_layers=3, last_bn=True) 

75 predictor = MLPHead(proj_dim, pred_hidden, proj_dim, num_layers=2, last_bn=False) 

76 return SimSiam(encoder, projector, predictor, anti_collapse=anti_collapse)