Ersi lig 3912 refactor mae to use timm vit

docs/source/getting_started/install.rst

@@ -33,7 +33,7 @@ If you want to work with video files you need to additionally install
     pip install av
 
 If you want to work use the Masked Autoencoder you need to additionally install
-`TIMM <https://timm.fast.ai/>`_.
+`TIMM <https://github.com/huggingface/pytorch-image-models>`_.
 
 .. code-block:: bash
 

examples/pytorch/mae.py

@@ -1,24 +1,13 @@
 # Note: The model and training settings do not follow the reference settings
 # from the paper. The settings are chosen such that the example can easily be
 # run on a small dataset with a single GPU.
-import sys
-
 import torch
 import torchvision
+from timm.models.vision_transformer import vit_base_patch32_224
 from torch import nn
 
-from lightly.utils import dependency
-
-if dependency.timm_vit_available():
-    from timm.models.vision_transformer import vit_base_patch32_224
-else:
-    sys.exit(1)
-
 from lightly.models import utils
-from lightly.models.modules import (
-    masked_autoencoder_timm,
-    masked_vision_transformer_timm,
-)
+from lightly.models.modules import MAEDecoderTIMM, MaskedVisionTransformerTIMM
 from lightly.transforms.mae_transform import MAETransform
 
 
@@ -30,16 +19,14 @@ def __init__(self, vit):
         self.mask_ratio = 0.75
         self.patch_size = vit.patch_embed.patch_size[0]
 
-        self.backbone = masked_vision_transformer_timm.MaskedVisionTransformerTIMM(
-            vit=vit
-        )
+        self.backbone = MaskedVisionTransformerTIMM(vit=vit)
         self.sequence_length = self.backbone.sequence_length
-        self.decoder = masked_autoencoder_timm.MAEDecoder(
+        self.decoder = MAEDecoderTIMM(
             num_patches=vit.patch_embed.num_patches,
             patch_size=self.patch_size,
             embed_dim=vit.embed_dim,
             decoder_embed_dim=decoder_dim,
-            decoder_depth=8,
+            decoder_depth=1,
             decoder_num_heads=16,
             mlp_ratio=4.0,
             proj_drop_rate=0.0,

examples/pytorch/msn.py

@@ -9,10 +9,8 @@
 
 from lightly.loss import MSNLoss
 from lightly.models import utils
+from lightly.models.modules import MaskedVisionTransformerTorchvision
 from lightly.models.modules.heads import MSNProjectionHead
-from lightly.models.modules.masked_vision_transformer_torchvision import (
-    MaskedVisionTransformerTorchvision,
-)
 from lightly.transforms.msn_transform import MSNTransform
 
 

examples/pytorch/pmsn.py

@@ -9,10 +9,8 @@
 
 from lightly.loss import PMSNLoss
 from lightly.models import utils
+from lightly.models.modules import MaskedVisionTransformerTorchvision
 from lightly.models.modules.heads import MSNProjectionHead
-from lightly.models.modules.masked_vision_transformer_torchvision import (
-    MaskedVisionTransformerTorchvision,
-)
 from lightly.transforms import MSNTransform
 
 

examples/pytorch_lightning/mae.py

@@ -1,25 +1,14 @@
 # Note: The model and training settings do not follow the reference settings
 # from the paper. The settings are chosen such that the example can easily be
 # run on a small dataset with a single GPU.
-import sys
-
 import pytorch_lightning as pl
 import torch
 import torchvision
+from timm.models.vision_transformer import vit_base_patch32_224
 from torch import nn
 
-from lightly.utils import dependency
-
-if dependency.timm_vit_available():
-    from timm.models.vision_transformer import vit_base_patch32_224
-else:
-    sys.exit(1)
-
 from lightly.models import utils
-from lightly.models.modules import (
-    masked_autoencoder_timm,
-    masked_vision_transformer_timm,
-)
+from lightly.models.modules import MAEDecoderTIMM, MaskedVisionTransformerTIMM
 from lightly.transforms.mae_transform import MAETransform
 
 
@@ -31,16 +20,14 @@ def __init__(self):
         vit = vit_base_patch32_224()
         self.mask_ratio = 0.75
         self.patch_size = vit.patch_embed.patch_size[0]
-        self.backbone = masked_vision_transformer_timm.MaskedVisionTransformerTIMM(
-            vit=vit
-        )
+        self.backbone = MaskedVisionTransformerTIMM(vit=vit)
         self.sequence_length = self.backbone.sequence_length
-        self.decoder = masked_autoencoder_timm.MAEDecoder(
+        self.decoder = MAEDecoderTIMM(
             num_patches=vit.patch_embed.num_patches,
             patch_size=self.patch_size,
             embed_dim=vit.embed_dim,
             decoder_embed_dim=decoder_dim,
-            decoder_depth=8,
+            decoder_depth=1,
             decoder_num_heads=16,
             mlp_ratio=4.0,
             proj_drop_rate=0.0,

examples/pytorch_lightning/msn.py

@@ -10,10 +10,8 @@
 
 from lightly.loss import MSNLoss
 from lightly.models import utils
+from lightly.models.modules import MaskedVisionTransformerTorchvision
 from lightly.models.modules.heads import MSNProjectionHead
-from lightly.models.modules.masked_vision_transformer_torchvision import (
-    MaskedVisionTransformerTorchvision,
-)
 from lightly.transforms.msn_transform import MSNTransform
 
 

examples/pytorch_lightning/pmsn.py

@@ -10,10 +10,8 @@
 
 from lightly.loss import PMSNLoss
 from lightly.models import utils
+from lightly.models.modules import MaskedVisionTransformerTorchvision
 from lightly.models.modules.heads import MSNProjectionHead
-from lightly.models.modules.masked_vision_transformer_torchvision import (
-    MaskedVisionTransformerTorchvision,
-)
 from lightly.transforms import MSNTransform
 
 

examples/pytorch_lightning/simmim.py

@@ -4,9 +4,7 @@
 from torch import nn
 
 from lightly.models import utils
-from lightly.models.modules.masked_vision_transformer_torchvision import (
-    MaskedVisionTransformerTorchvision,
-)
+from lightly.models.modules import MaskedVisionTransformerTorchvision
 from lightly.transforms.mae_transform import MAETransform  # Same transform as MAE
 
 

examples/pytorch_lightning_distributed/mae.py

@@ -1,25 +1,14 @@
 # Note: The model and training settings do not follow the reference settings
 # from the paper. The settings are chosen such that the example can easily be
 # run on a small dataset with a single GPU.
-import sys
-
 import pytorch_lightning as pl
 import torch
 import torchvision
+from timm.models.vision_transformer import vit_base_patch32_224
 from torch import nn
 
-from lightly.utils import dependency
-
-if dependency.timm_vit_available():
-    from timm.models.vision_transformer import vit_base_patch32_224
-else:
-    sys.exit(1)
-
 from lightly.models import utils
-from lightly.models.modules import (
-    masked_autoencoder_timm,
-    masked_vision_transformer_timm,
-)
+from lightly.models.modules import MAEDecoderTIMM, MaskedVisionTransformerTIMM
 from lightly.transforms.mae_transform import MAETransform
 
 
@@ -31,16 +20,14 @@ def __init__(self):
         vit = vit_base_patch32_224()
         self.mask_ratio = 0.75
         self.patch_size = vit.patch_embed.patch_size[0]
-        self.backbone = masked_vision_transformer_timm.MaskedVisionTransformerTIMM(
-            vit=vit
-        )
+        self.backbone = MaskedVisionTransformerTIMM(vit=vit)
         self.sequence_length = self.backbone.sequence_length
-        self.decoder = masked_autoencoder_timm.MAEDecoder(
+        self.decoder = MAEDecoderTIMM(
             num_patches=vit.patch_embed.num_patches,
             patch_size=self.patch_size,
             embed_dim=vit.embed_dim,
             decoder_embed_dim=decoder_dim,
-            decoder_depth=8,
+            decoder_depth=1,
             decoder_num_heads=16,
             mlp_ratio=4.0,
             proj_drop_rate=0.0,

examples/pytorch_lightning_distributed/msn.py

@@ -10,10 +10,8 @@
 
 from lightly.loss import MSNLoss
 from lightly.models import utils
+from lightly.models.modules import MaskedVisionTransformerTorchvision
 from lightly.models.modules.heads import MSNProjectionHead
-from lightly.models.modules.masked_vision_transformer_torchvision import (
-    MaskedVisionTransformerTorchvision,
-)
 from lightly.transforms.msn_transform import MSNTransform
 
 
@@ -117,7 +115,6 @@ def configure_optimizers(self):
     devices="auto",
     accelerator="gpu",
     strategy="ddp",
-    # use_distributed_sampler=True,  # or replace_sampler_ddp=True for PyTorch Lightning <2.0
-    replace_sampler_ddp=True,
+    use_distributed_sampler=True,  # or replace_sampler_ddp=True for PyTorch Lightning <2.0
 )
 trainer.fit(model=model, train_dataloaders=dataloader)

examples/pytorch_lightning_distributed/pmsn.py

@@ -10,10 +10,8 @@
 
 from lightly.loss import PMSNLoss
 from lightly.models import utils
+from lightly.models.modules import MaskedVisionTransformerTorchvision
 from lightly.models.modules.heads import MSNProjectionHead
-from lightly.models.modules.masked_vision_transformer_torchvision import (
-    MaskedVisionTransformerTorchvision,
-)
 from lightly.transforms import MSNTransform
 
 
@@ -118,7 +116,6 @@ def configure_optimizers(self):
     devices="auto",
     accelerator="gpu",
     strategy="ddp",
-    # use_distributed_sampler=True,  # or replace_sampler_ddp=True for PyTorch Lightning <2.0
-    replace_sampler_ddp=True,
+    use_distributed_sampler=True,  # or replace_sampler_ddp=True for PyTorch Lightning <2.0
 )
 trainer.fit(model=model, train_dataloaders=dataloader)

examples/pytorch_lightning_distributed/simmim.py

@@ -4,9 +4,7 @@
 from torch import nn
 
 from lightly.models import utils
-from lightly.models.modules.masked_vision_transformer_torchvision import (
-    MaskedVisionTransformerTorchvision,
-)
+from lightly.models.modules import MaskedVisionTransformerTorchvision
 from lightly.transforms.mae_transform import MAETransform  # Same transform as MAE
 
 
@@ -97,7 +95,6 @@ def configure_optimizers(self):
     devices="auto",
     accelerator="gpu",
     strategy="ddp",
-    # use_distributed_sampler=True,  # or replace_sampler_ddp=True for PyTorch Lightning <2.0
-    replace_sampler_ddp=True,
+    use_distributed_sampler=True,  # or replace_sampler_ddp=True for PyTorch Lightning <2.0
 )
 trainer.fit(model=model, train_dataloaders=dataloader)

lightly/models/modules/__init__.py

@@ -31,7 +31,6 @@
 from lightly.utils import dependency as _dependency
 
 if _dependency.torchvision_vit_available():
-    # Requires torchvision >=0.12
     # Requires torchvision >=0.12
     from lightly.models.modules.masked_autoencoder import (
         MAEBackbone,
@@ -45,10 +44,10 @@
 if _dependency.timm_vit_available():
     # Requires timm >= 0.9.9
     from lightly.models.modules.heads_timm import AIMPredictionHead
-    from lightly.models.modules.masked_autoencoder_timm import MAEDecoder
+    from lightly.models.modules.masked_autoencoder_timm import MAEDecoderTIMM
     from lightly.models.modules.masked_causal_vision_transformer import (
         MaskedCausalVisionTransformer,
     )
     from lightly.models.modules.masked_vision_transformer_timm import (
         MaskedVisionTransformerTIMM,
     )

lightly/models/modules/masked_autoencoder_timm.py

@@ -1,22 +1,17 @@
 from __future__ import annotations
 
 from functools import partial
 from typing import Callable, Optional
 
-from lightly.utils import dependency
-
-if dependency.timm_vit_available():
-    from timm.models import vision_transformer
-
-
 import torch
 import torch.nn as nn
+from timm.models import vision_transformer
 from torch.nn import LayerNorm, Linear, Module, Parameter, Sequential
 
 from lightly.models import utils
 
 
-class MAEDecoder(Module):
+class MAEDecoderTIMM(Module):
     """Decoder for the Masked Autoencoder model [0].
 
     Decodes encoded patches and predicts pixel values for every patch.
@@ -53,7 +48,7 @@
 
     """
 
     def __init__(
         self,
         num_patches: int,
         patch_size: int,
@@ -68,21 +63,21 @@
         norm_layer: Callable[..., nn.Module] = partial(LayerNorm, eps=1e-6),
         mask_token: Optional[Parameter] = None,
     ):
         super().__init__()
 
         self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
         self.mask_token = (
             nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
             if mask_token is None
             else mask_token
         )
 
         # positional encoding of the decoder
         self.decoder_pos_embed = nn.Parameter(
             torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False
         )  # fixed sin-cos embedding
 
         self.decoder_blocks = Sequential(
             *[
                 vision_transformer.Block(
                     decoder_embed_dim,
@@ -97,14 +92,14 @@
             ]
         )
 
         self.decoder_norm = norm_layer(decoder_embed_dim)
         self.decoder_pred = nn.Linear(
             decoder_embed_dim, patch_size**2 * in_chans, bias=True
         )  # decoder to patch
 
         self._initialize_weights()
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         """Returns predicted pixel values from encoded tokens.
 
         Args:
@@ -115,11 +110,11 @@
             Tensor with shape (batch_size, seq_length, out_dim).
 
         """
         out = self.embed(input)
         out = self.decode(out)
         return self.predict(out)
 
     def embed(self, input: torch.Tensor) -> torch.Tensor:
         """Embeds encoded input tokens into decoder token dimension.
 
         This is a single linear layer that changes the token dimension from
@@ -135,10 +130,10 @@
             the embedded tokens.
 
         """
         out: torch.Tensor = self.decoder_embed(input)
         return out
 
     def decode(self, input: torch.Tensor) -> torch.Tensor:
         """Forward pass through the decoder transformer.
 
         Args:
@@ -151,12 +146,12 @@
             the decoded tokens.
 
         """
         output: torch.Tensor = input + self.decoder_pos_embed
         output = self.decoder_blocks(output)
         output = self.decoder_norm(output)
         return output
 
     def predict(self, input: torch.Tensor) -> torch.Tensor:
         """Predics pixel values from decoded tokens.
 
         Args:
@@ -169,35 +164,35 @@
             predictions for each token.
 
         """
         out: torch.Tensor = self.decoder_pred(input)
         return out
 
     def _initialize_weights(self) -> None:
         torch.nn.init.normal_(self.mask_token, std=0.02)
         _initialize_2d_sine_cosine_positional_embedding(self.decoder_pos_embed)
         self.apply(_init_weights)
 
 
 def _initialize_2d_sine_cosine_positional_embedding(pos_embedding: Parameter) -> None:
     _, seq_length, hidden_dim = pos_embedding.shape
     grid_size = int((seq_length - 1) ** 0.5)
     sine_cosine_embedding = utils.get_2d_sine_cosine_positional_embedding(
         embed_dim=hidden_dim,
         grid_size=grid_size,
         cls_token=True,
     )
     pos_embedding.data.copy_(
         torch.from_numpy(sine_cosine_embedding).float().unsqueeze(0)
     )
     # Freeze positional embedding.
     pos_embedding.requires_grad = False
 
 
 def _init_weights(module: Module) -> None:
     if isinstance(module, Linear):
         nn.init.xavier_uniform_(module.weight)
         if isinstance(module, Linear) and module.bias is not None:
             nn.init.constant_(module.bias, 0)
     elif isinstance(module, LayerNorm):
         nn.init.constant_(module.bias, 0)
         nn.init.constant_(module.weight, 1.0)

lightly/utils/dependency.py

@@ -21,7 +21,6 @@ def timm_vit_available() -> bool:
         import timm.models.vision_transformer  # Requires timm >= 0.3.3
         from timm.layers import LayerType  # Requires timm >= 0.9.9
     except ImportError:
-        print("TIMM is not available. Please install if you would like to use the MAE.")
         return False
     else:
         return True