Add GLIP (Centered Masking for Language-Image Pre-Training)

src/open_clip/factory.py

@@ -201,6 +201,8 @@ def create_model(
         force_quick_gelu: bool = False,
         force_custom_text: bool = False,
         force_patch_dropout: Optional[float] = None,
+        gaussian_masking: Optional[bool] = None,
+        gaussian_masking_std: Optional[float] = None,
         force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
         force_preprocess_cfg: Optional[Dict[str, Any]] = None,
         pretrained_image: bool = False,
@@ -251,6 +253,14 @@ def create_model(
         if force_patch_dropout is not None:
             # override the default patch dropout value
             model_cfg["vision_cfg"]["patch_dropout"] = force_patch_dropout
+
+        if gaussian_masking is not None:
+            # override the default gaussian masking value
+            model_cfg["vision_cfg"]["gaussian_masking"] = gaussian_masking
+
+        if gaussian_masking_std is not None:
+            # override the default gaussian masking std value
+            model_cfg["vision_cfg"]["gaussian_masking_std"] = gaussian_masking_std
 
         if force_image_size is not None:
             # override model config's image size
@@ -396,6 +406,8 @@ def create_model_and_transforms(
         force_quick_gelu: bool = False,
         force_custom_text: bool = False,
         force_patch_dropout: Optional[float] = None,
+        gaussian_masking: Optional[bool] = None,
+        gaussian_masking_std: Optional[float] = None,
         force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
         image_mean: Optional[Tuple[float, ...]] = None,
         image_std: Optional[Tuple[float, ...]] = None,
@@ -420,6 +432,8 @@ def create_model_and_transforms(
         force_quick_gelu=force_quick_gelu,
         force_custom_text=force_custom_text,
         force_patch_dropout=force_patch_dropout,
+        gaussian_masking=gaussian_masking,
+        gaussian_masking_std=gaussian_masking_std,
         force_image_size=force_image_size,
         force_preprocess_cfg=force_preprocess_cfg,
         pretrained_image=pretrained_image,

src/open_clip/model.py

@@ -34,6 +34,8 @@ class CLIPVisionCfg:
 
     ls_init_value: Optional[float] = None  # layer scale initial value
     patch_dropout: float = 0.  # what fraction of patches to dropout during training (0 would mean disabled and no patches dropped) - 0.5 to 0.75 recommended in the paper for optimal results
+    gaussian_masking: bool = False
+    gaussian_masking_std: float = 0.20
     attentional_pool: bool = False  # whether to use attentional pooler in the last embedding layer (overrides pool_type)
     attn_pooler_queries: int = 256  # n_queries for attentional pooler
     attn_pooler_heads: int = 8  # n heads for attentional_pooling
@@ -155,6 +157,8 @@ def _build_vision_tower(
             mlp_ratio=vision_cfg.mlp_ratio,
             ls_init_value=vision_cfg.ls_init_value,
             patch_dropout=vision_cfg.patch_dropout,
+            gaussian_masking=vision_cfg.gaussian_masking,
+            gaussian_masking_std=vision_cfg.gaussian_masking_std,
             attentional_pool=vision_cfg.attentional_pool,
             attn_pooler_queries=vision_cfg.attn_pooler_queries,
             attn_pooler_heads=vision_cfg.attn_pooler_heads,

src/open_clip/transformer.py

@@ -7,6 +7,8 @@
 from torch import nn
 from torch.nn import functional as F
 from torch.utils.checkpoint import checkpoint
+from scipy.stats import multivariate_normal
+import numpy as np
 
 from .utils import to_2tuple
 from .pos_embed import get_2d_sincos_pos_embed
@@ -51,11 +53,29 @@ class PatchDropout(nn.Module):
     https://arxiv.org/abs/2212.00794
     """
 
-    def __init__(self, prob, exclude_first_token=True):
+    def __init__(self, prob, grid_size, gaussian_masking=False, std=0.20, exclude_first_token=True):
         super().__init__()
         assert 0 <= prob < 1.
         self.prob = prob
         self.exclude_first_token = exclude_first_token  # exclude CLS token
+        self.gaussian_masking = gaussian_masking        
+        self.pdf_values = self.normal_distribution(grid_size, std=std)
+
+
+    def normal_distribution(self, grid_size=14, mean=0, std=0.20):
+        """
+        https://arxiv.org/abs/2403.15837
+        """
+
+        x = np.linspace(-1, 1, grid_size)
+        x, y = np.meshgrid(x, x)
+        mean, cov_matrix = [mean, mean], [[std, 0], [0, std]]
+        bivariate_dist = multivariate_normal(mean=mean, cov=cov_matrix)
+        points = np.column_stack([x.flatten(), y.flatten()])
+        # Evaluate the PDF at the given points
+        pdf_values = bivariate_dist.pdf(points)
+
+        return 1 - pdf_values
 
     def forward(self, x):
         if not self.training or self.prob == 0.:
@@ -76,6 +96,8 @@ def forward(self, x):
         num_patches_keep = max(1, int(num_tokens * keep_prob))
 
         rand = torch.randn(batch, num_tokens)
+        if self.gaussian_masking:
+            rand = rand - self.pdf_values
         patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
 
         x = x[batch_indices, patch_indices_keep]
@@ -448,6 +470,8 @@ def __init__(
             attn_pooler_heads: int = 8,
             output_dim: int = 512,
             patch_dropout: float = 0.,
+            gaussian_masking: bool = False,
+            gaussian_masking_std: float = 0.20,
             no_ln_pre: bool = False,
             pos_embed_type: str = 'learnable',
             pool_type: str = 'tok',
@@ -485,7 +509,8 @@ def __init__(
             raise ValueError
 
         # setting a patch_dropout of 0. would mean it is disabled and this function would be the identity fn
-        self.patch_dropout = PatchDropout(patch_dropout) if patch_dropout > 0. else nn.Identity()
+        self.patch_dropout = PatchDropout(patch_dropout, self.grid_size[0],
+                                          gaussian_masking, gaussian_masking_std) if patch_dropout > 0. else nn.Identity()
 
         self.ln_pre = nn.Identity() if no_ln_pre else norm_layer(width)
         self.transformer = Transformer(

src/open_clip_train/main.py

@@ -229,6 +229,8 @@ def main(args):
         force_quick_gelu=args.force_quick_gelu,
         force_custom_text=args.force_custom_text,
         force_patch_dropout=args.force_patch_dropout,
+        gaussian_masking=args.gaussian_masking,
+        gaussian_masking_std=args.gaussian_masking_std,
         force_image_size=args.force_image_size,
         image_mean=args.image_mean,
         image_std=args.image_std,

src/open_clip_train/params.py

@@ -279,6 +279,18 @@ def parse_args(args):
         type=float,
         help="Override the patch dropout during training, for fine tuning with no dropout near the end as in the paper",
     )
+    parser.add_argument(
+        "--gaussian-masking",
+        default=False,
+        action='store_true',
+        help="Use gaussian masking for patch dropout.",
+    )
+    parser.add_argument(
+        "--gaussian-masking-std",
+        default=0.20,
+        type=float,
+        help="Gaussian masking std for patch dropout.",
+    )
     parser.add_argument(
         "--force-custom-text",
         default=False,