diff --git a/lightly/models/modules/ijepa.py b/lightly/models/modules/ijepa.py index 7889dadd1..f112efa3d 100644 --- a/lightly/models/modules/ijepa.py +++ b/lightly/models/modules/ijepa.py @@ -79,6 +79,11 @@ def __init__( torch.from_numpy(predictor_pos_embed).float().unsqueeze(0) ) + self.use_stop = kwargs.get( + "use_stop", False + ) # pass use stop embeddings as additional args, default to False + self.noise_std = kwargs.get("noise_std", 0.25) # default 0.25 + @classmethod def from_vit_encoder(cls, vit_encoder, num_patches): """Creates an I-JEPA predictor backbone (multi-head attention and layernorm) from a torchvision ViT encoder. @@ -134,6 +139,7 @@ def forward(self, x, masks_x, masks): if not isinstance(masks, list): masks = [masks] + noise_dim = x.shape[-1] B = len(x) // len(masks_x) x = self.predictor_embed(x) x_pos_embed = self.predictor_pos_embed.repeat(B, 1, 1) @@ -144,9 +150,21 @@ def forward(self, x, masks_x, masks): pos_embs = self.predictor_pos_embed.repeat(B, 1, 1) pos_embs = utils.apply_masks(pos_embs, masks) pos_embs = utils.repeat_interleave_batch(pos_embs, B, repeat=len(masks_x)) + + # we add the stochastic positional embedding here: + # use self.predictor_embed as the projector + pos_embs = utils.add_stochastic_positional_noise( + pos_embs, + self.predictor_embed, + noise_dim, + noise_std=self.noise_std, + enabled=self.use_stop, + ) + pred_tokens = self.mask_token.repeat(pos_embs.size(0), pos_embs.size(1), 1) pred_tokens += pos_embs + x = x.repeat(len(masks), 1, 1) x = torch.cat([x, pred_tokens], dim=1) diff --git a/lightly/models/modules/ijepa_timm.py b/lightly/models/modules/ijepa_timm.py index ebba789ab..adfe57d16 100644 --- a/lightly/models/modules/ijepa_timm.py +++ b/lightly/models/modules/ijepa_timm.py @@ -61,6 +61,8 @@ def __init__( proj_drop_rate: float = 0.0, attn_drop_rate: float = 0.0, norm_layer: Callable[..., nn.Module] = partial(nn.LayerNorm, eps=1e-6), + use_stop: bool = False, + noise_std: float = 0.25, ): """Initializes the IJEPAPredictorTIMM with the specified dimensions.""" super().__init__() @@ -97,6 +99,9 @@ def __init__( ] ) + self.use_stop = use_stop + self.noise_std = noise_std + def forward( self, x: Tensor, @@ -123,6 +128,7 @@ def forward( len_masks_x = len(masks_x) if isinstance(masks_x, list) else 1 len_masks = len(masks) if isinstance(masks, list) else 1 + noise_dim = x.shape[-1] B = len(x) // len_masks_x x = self.predictor_embed(x) x_pos_embed = self.predictor_pos_embed.repeat(B, 1, 1) @@ -136,6 +142,17 @@ def forward( pred_tokens = self.mask_token.repeat(pos_embs.size(0), pos_embs.size(1), 1) pred_tokens += pos_embs + + # we add the stochastic positional embedding here: + # use self.predictor_embed as the projector + pred_tokens = utils.add_stochastic_positional_noise( + pred_tokens, + self.predictor_embed, + noise_dim, + noise_std=self.noise_std, + enabled=self.use_stop, + ) + x = x.repeat(len_masks, 1, 1) x = torch.cat([x, pred_tokens], dim=1) diff --git a/lightly/models/utils.py b/lightly/models/utils.py index eb6c17533..31d3069f3 100644 --- a/lightly/models/utils.py +++ b/lightly/models/utils.py @@ -1316,3 +1316,46 @@ def apply_masks(x: Tensor, masks: Tensor | list[Tensor]) -> Tensor: mask_keep = m.unsqueeze(-1).repeat(1, 1, x.size(-1)) all_x += [torch.gather(x, dim=1, index=mask_keep)] return torch.cat(all_x, dim=0) + + +def add_stochastic_positional_noise( + pos_embeddings: Tensor, + projection: Module, + noise_dim: int, + noise_std: float = 0.25, + enabled: bool = False, +) -> Tensor: + """Adds stochastic noise to positional embeddings. + + [0]. https://arxiv.org/pdf/2308.00566 + [1]. https://github.com/amirbar/StoP/blob/main/src/deit.py + + Args: + pos_embeddings: + Positional embeddings of shape + ``(batch_size, num_tokens, predictor_embed_dim)``. + projection: + Matrix A used to project gaussian noise to the pos_embedding + dimension. + noise_dim: + Dimension of the sampled gaussian noise before projection. + noise_std: + Standard deviation of the gaussian noise. + enabled: + If False, returns ``pos_embeddings`` unchanged. + + Returns: + Positional embeddings with optional gaussian noise added. + """ + if not enabled or noise_std == 0.0: + return pos_embeddings + + noise = torch.normal( + mean=0.0, + std=noise_std, + size=(pos_embeddings.shape[0], pos_embeddings.shape[1], noise_dim), + device=pos_embeddings.device, + dtype=pos_embeddings.dtype, + ) + + return pos_embeddings + projection(noise) diff --git a/tests/models/modules/test_ijepa_timm.py b/tests/models/modules/test_ijepa_timm.py index f41a489cf..6cc1b9ccf 100644 --- a/tests/models/modules/test_ijepa_timm.py +++ b/tests/models/modules/test_ijepa_timm.py @@ -14,8 +14,10 @@ from lightly.models.modules import IJEPAPredictorTIMM -class TestIJEPAPredictorTIMM(unittest.TestCase): - def test_init(self) -> None: +class TestIJEPAPredictorTIMM: + @pytest.mark.parametrize("use_stop", [True, False]) + @pytest.mark.parametrize("noise_std", [0.0, 0.1]) + def test_init(self, use_stop: bool, noise_std: float) -> None: IJEPAPredictorTIMM( num_patches=196, depth=2, @@ -26,10 +28,17 @@ def test_init(self) -> None: mlp_ratio=4.0, proj_drop_rate=0.0, attn_drop_rate=0.0, + use_stop=use_stop, + noise_std=noise_std, ) def _test_forward( - self, device: torch.device, batch_size: int = 4, seed: int = 0 + self, + device: torch.device, + use_stop: bool, + noise_std: float, + batch_size: int = 4, + seed: int = 0, ) -> None: torch.manual_seed(seed) num_patches = 196 # 14x14 patches @@ -48,6 +57,8 @@ def _test_forward( mlp_ratio=4.0, proj_drop_rate=0.0, attn_drop_rate=0.0, + use_stop=use_stop, + noise_std=noise_std, ).to(device) x = torch.randn(batch_size, num_patches, mlp_dim, device=device) @@ -56,16 +67,16 @@ def _test_forward( predictions = predictor(x, masks_x, masks) - # output shape must be correct - expected_shape = [batch_size, num_patches, mlp_dim] - self.assertListEqual(list(predictions.shape), expected_shape) + assert list(predictions.shape) == [batch_size, num_patches, mlp_dim] + assert torch.all(torch.isfinite(predictions)) - # output must have reasonable numbers - self.assertTrue(torch.all(torch.isfinite(predictions))) + @pytest.mark.parametrize("use_stop", [True, False]) + @pytest.mark.parametrize("noise_std", [0.0, 0.1]) + def test_forward(self, use_stop: bool, noise_std: float) -> None: + self._test_forward(torch.device("cpu"), use_stop, noise_std) - def test_forward(self) -> None: - self._test_forward(torch.device("cpu")) - - @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available.") - def test_forward_cuda(self) -> None: - self._test_forward(torch.device("cuda")) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available.") + @pytest.mark.parametrize("use_stop", [True, False]) + @pytest.mark.parametrize("noise_std", [0.0, 0.1]) + def test_forward_cuda(self, use_stop: bool, noise_std: float) -> None: + self._test_forward(torch.device("cuda"), use_stop, noise_std) diff --git a/tests/utils/test_stochastic_positional_embedding.py b/tests/utils/test_stochastic_positional_embedding.py new file mode 100644 index 000000000..19bfe0958 --- /dev/null +++ b/tests/utils/test_stochastic_positional_embedding.py @@ -0,0 +1,31 @@ +import torch + +from lightly.models import utils + + +def test_add_stochastic_positional_noise_disabled() -> None: + projection = torch.nn.Linear(8, 4) + pos_embeddings = torch.randn(2, 3, 4) + + out = utils.add_stochastic_positional_noise( + pos_embeddings=pos_embeddings, + projection=projection, + noise_dim=8, + enabled=False, + ) + + assert torch.equal(out, pos_embeddings) + + +def test_add_stochastic_positional_noise_enabled_shape() -> None: + projection = torch.nn.Linear(8, 4) + pos_embeddings = torch.randn(2, 3, 4) + + out = utils.add_stochastic_positional_noise( + pos_embeddings=pos_embeddings, + projection=projection, + noise_dim=8, + enabled=True, + ) + + assert out.shape == pos_embeddings.shape