Add TextTextCLIP

docs/script_examples/train_clanp.sh

@@ -0,0 +1,62 @@
+#!/bin/bash
+#SBATCH --partition=g40
+#SBATCH --account=laion
+#SBATCH --job-name=CLANP
+#SBATCH --nodes=4
+#SBATCH --gpus-per-node=8
+#SBATCH --cpus-per-gpu=10
+#SBATCH --output=logs/%x_%j.out
+#SBATCH --comment=laion
+#SBATCH --open-mode=append
+#SBATCH --exclusive
+source /etc/profile.d/modules.sh
+module load openmpi
+module load cuda/11.7
+
+export PYTHONFAULTHANDLER=1
+export CUDA_LAUNCH_BLOCKING=0
+export HOSTNAMES=`scontrol show hostnames "$SLURM_JOB_NODELIST"`
+export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+export MASTER_PORT=12802
+export COUNT_NODE=`scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l`
+
+echo go $COUNT_NODE
+echo $HOSTNAMES
+
+
+cd /admin/home-jianz/open_clip/src
+export PYTHONPATH="$PYTHONPATH:/admin/home-jianz/open_clip/src"
+
+EXP_NAME=""
+
+
+
+srun --comment laion --cpu_bind=v --accel-bind=gn python3 -m training.main \
+    --save-frequency 1 \
+    --dataset-type="webdataset" \
+    --text-a-key="txt" \
+    --text-b-key="txt" \
+    --report-to wandb \
+    --wandb-project-name="CLANP" \
+    --train-data="pipe:aws s3 --cli-connect-timeout 1000 --cli-read-timeout 1000 cp s3://s-laion/pile/pile_wds/pile-{00000..01159}.tar -" \
+    --train-num-samples=1000000 \
+    --warmup 5000 \
+    --batch-size=256 \
+    --lr=0.00001 \
+    --wd=0.01 \
+    --epochs=100 \
+    --workers=10 \
+    --model="Siamese-xlm-roberta-large" \
+    --save-most-recent \
+    --seed 42 \
+    --local-loss \
+    --gather-with-grad \
+    --ddp-static-graph \
+    --grad-checkpointing \
+    --precision amp_bfloat16 \
+    --model-type "SiameseCLANP" \
+    --unsupervised-pretraining \
+    --context-length=256 \
+    --logs /fsx/home-jianz/logs \
+    --resume /fsx/home-jianz/logs/2023_05_25-09_21_35-model_Siamese-xlm-roberta-large-lr_1e-05-b_256-j_10-p_amp_bfloat16/checkpoints/epoch_latest.pt
+

requirements-training.txt

@@ -8,5 +8,7 @@ pandas
 braceexpand
 huggingface_hub
 transformers
+datasets
 timm
 fsspec
+scipy

src/open_clip/__init__.py

@@ -2,8 +2,9 @@
 from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
 from .factory import create_model, create_model_and_transforms, create_model_from_pretrained, get_tokenizer, create_loss
 from .factory import list_models, add_model_config, get_model_config, load_checkpoint
+
 from .loss import ClipLoss, DistillClipLoss, CoCaLoss
-from .model import CLIP, CustomTextCLIP, CLIPTextCfg, CLIPVisionCfg, \
+from .model import CLIP, CustomTextCLIP, SiameseCLANP, CLANP, CLIPTextCfg, CLIPVisionCfg,\
     convert_weights_to_lp, convert_weights_to_fp16, trace_model, get_cast_dtype
 from .openai import load_openai_model, list_openai_models
 from .pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model, \

src/open_clip/factory.py

@@ -10,7 +10,7 @@
 import torch
 
 from .constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
-from .model import CLIP, CustomTextCLIP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
+from .model import CLIP, CustomTextCLIP, CLANP, SiameseCLANP, convert_weights_to_lp, convert_to_custom_text_state_dict,\
     resize_pos_embed, get_cast_dtype
 from .coca_model import CoCa
 from .loss import ClipLoss, DistillClipLoss, CoCaLoss
@@ -44,10 +44,16 @@ def _rescan_model_configs():
     for cf in config_files:
         with open(cf, 'r') as f:
             model_cfg = json.load(f)
-            if all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg')):
+
+            is_clip = all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg'))
+            is_text_only = all(a in model_cfg for a in ('embed_dim', 'tower_a_cfg', 'tower_b_cfg'))
+            is_siamese_text_only = all(a in model_cfg for a in ('embed_dim', 'text_cfg'))
+            is_multimodal = all(a in model_cfg for a in ('embed_dim', 'vision_cfg', 'text_cfg', 'multimodal_cfg'))
+            if is_clip or is_text_only or is_siamese_text_only or is_multimodal:
                 _MODEL_CONFIGS[cf.stem] = model_cfg
 
     _MODEL_CONFIGS = {k: v for k, v in sorted(_MODEL_CONFIGS.items(), key=lambda x: _natural_key(x[0]))}
+
 
 
 _rescan_model_configs()  # initial populate of model config registry
@@ -74,12 +80,24 @@ def get_model_config(model_name):
 
 
 def get_tokenizer(model_name):
+
     if model_name.startswith(HF_HUB_PREFIX):
         tokenizer = HFTokenizer(model_name[len(HF_HUB_PREFIX):])
     else:
         config = get_model_config(model_name)
-        tokenizer = HFTokenizer(
-            config['text_cfg']['hf_tokenizer_name']) if 'hf_tokenizer_name' in config['text_cfg'] else tokenize
+
+        if 'text_cfg' in config.keys():
+            key = 'text_cfg'
+        elif 'tower_a_cfg' in config.keys():
+            key = 'tower_a_cfg'
+        if 'hf_tokenizer_name' in config[key]:
+            tokenizer = HFTokenizer(config[key]['hf_tokenizer_name']) 
+            if "pythia" in config[key]['hf_tokenizer_name']:
+                tokenizer.tokenizer.pad_token_id = 1
+        else: 
+            tokenizer = tokenize
+
+
     return tokenizer
 
 
@@ -117,8 +135,10 @@ def create_model(
         pretrained_image: bool = False,
         pretrained_hf: bool = True,
         cache_dir: Optional[str] = None,
+        model_type: Optional[str] = "CLIP",
         output_dict: Optional[bool] = None,
         require_pretrained: bool = False,
+
 ):
     has_hf_hub_prefix = model_name.startswith(HF_HUB_PREFIX)
     if has_hf_hub_prefix:
@@ -180,18 +200,35 @@ def create_model(
                 assert False, 'pretrained image towers currently only supported for timm models'
 
         cast_dtype = get_cast_dtype(precision)
+
         is_hf_model = 'hf_model_name' in model_cfg.get('text_cfg', {})
         custom_text = model_cfg.pop('custom_text', False) or force_custom_text or is_hf_model
 
-        if custom_text:
-            if is_hf_model:
-                model_cfg['text_cfg']['hf_model_pretrained'] = pretrained_hf
-            if "coca" in model_name:
-                model = CoCa(**model_cfg, cast_dtype=cast_dtype)
+        # switch to CLANP
+        if model_type=="CLANP":
+            if 'hf_model_name' in model_cfg.get('tower_a_cfg', {}):
+                model_cfg['tower_a_cfg']['hf_model_pretrained'] = pretrained_hf
+            if 'hf_model_name' in model_cfg.get('tower_b_cfg', {}):
+                model_cfg['tower_b_cfg']['hf_model_pretrained'] = pretrained_hf
+            model = CLANP(**model_cfg, cast_dtype=cast_dtype)
+
+        elif model_type=="SiameseCLANP":
+            if 'hf_model_name' in model_cfg.get('tower_a_cfg', {}):
+                model_cfg['tower_a_cfg']['hf_model_pretrained'] = pretrained_hf
+            model = SiameseCLANP(**model_cfg, cast_dtype=cast_dtype)
+
+        elif model_type=="CLIP":
+
+            if custom_text:
+                if is_hf_model:
+                    model_cfg['text_cfg']['hf_model_pretrained'] = pretrained_hf
+                if "coca" in model_name:
+                    model = CoCa(**model_cfg, cast_dtype=cast_dtype)
+                else:
+                    model = CustomTextCLIP(**model_cfg, cast_dtype=cast_dtype)
             else:
-                model = CustomTextCLIP(**model_cfg, cast_dtype=cast_dtype)
-        else:
-            model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+                model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+
 
         pretrained_loaded = False
         if pretrained:
@@ -226,9 +263,10 @@ def create_model(
         if precision in ("fp16", "bf16"):
             convert_weights_to_lp(model, dtype=torch.bfloat16 if precision == 'bf16' else torch.float16)
 
+        if model_type=="CLIP":
         # set image / mean metadata from pretrained_cfg if available, or use default
-        model.visual.image_mean = pretrained_cfg.get('mean', None) or OPENAI_DATASET_MEAN
-        model.visual.image_std = pretrained_cfg.get('std', None) or OPENAI_DATASET_STD
+            model.visual.image_mean = pretrained_cfg.get('mean', None) or OPENAI_DATASET_MEAN
+            model.visual.image_std = pretrained_cfg.get('std', None) or OPENAI_DATASET_STD
 
         # to always output dict even if it is clip
         if output_dict and hasattr(model, "output_dict"):
@@ -287,6 +325,7 @@ def create_model_and_transforms(
         image_std: Optional[Tuple[float, ...]] = None,
         aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
         cache_dir: Optional[str] = None,
+        model_type: Optional[str] = "CLIP",
         output_dict: Optional[bool] = None,
 ):
     model = create_model(
@@ -302,24 +341,29 @@ def create_model_and_transforms(
         pretrained_image=pretrained_image,
         pretrained_hf=pretrained_hf,
         cache_dir=cache_dir,
+        model_type=model_type,
         output_dict=output_dict,
     )
 
-    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
-    image_std = image_std or getattr(model.visual, 'image_std', None)
-    preprocess_train = image_transform(
-        model.visual.image_size,
-        is_train=True,
-        mean=image_mean,
-        std=image_std,
-        aug_cfg=aug_cfg,
-    )
-    preprocess_val = image_transform(
-        model.visual.image_size,
-        is_train=False,
-        mean=image_mean,
-        std=image_std,
-    )
+    if model_type=="CLIP":
+        image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+        image_std = image_std or getattr(model.visual, 'image_std', None)
+        preprocess_train = image_transform(
+            model.visual.image_size,
+            is_train=True,
+            mean=image_mean,
+            std=image_std
+        )
+        preprocess_val = image_transform(
+            model.visual.image_size,
+            is_train=False,
+            mean=image_mean,
+            std=image_std
+        )
+    else:
+        preprocess_val = None
+        preprocess_train = None
+
 
     return model, preprocess_train, preprocess_val
 
@@ -337,6 +381,7 @@ def create_model_from_pretrained(
         image_mean: Optional[Tuple[float, ...]] = None,
         image_std: Optional[Tuple[float, ...]] = None,
         cache_dir: Optional[str] = None,
+        model_type: Optional[str] = "CLIP",
 ):
     model = create_model(
         model_name,
@@ -348,19 +393,23 @@ def create_model_from_pretrained(
         force_custom_text=force_custom_text,
         force_image_size=force_image_size,
         cache_dir=cache_dir,
+        model_type=model_type,
         require_pretrained=True,
     )
 
     if not return_transform:
         return model
 
-    image_mean = image_mean or getattr(model.visual, 'image_mean', None)
-    image_std = image_std or getattr(model.visual, 'image_std', None)
-    preprocess = image_transform(
-        model.visual.image_size,
-        is_train=False,
-        mean=image_mean,
-        std=image_std,
-    )
+
+    if model_type=="CLIP":
+        image_mean = image_mean or getattr(model.visual, 'image_mean', None)
+        image_std = image_std or getattr(model.visual, 'image_std', None)
+        preprocess = image_transform(
+            model.visual.image_size,
+            is_train=False,
+            mean=image_mean,
+            std=image_std
+        )
+
 
     return model, preprocess

src/open_clip/hf_configs.py

@@ -42,4 +42,18 @@
         },
         "pooler": "mean_pooler",
     },
+    # https://huggingface.co/docs/transformers/main/en/model_doc/gpt_neox#transformers.GPTNeoXConfig
+    "gpt_neox": {
+        "config_names": {
+            # https://github.com/huggingface/transformers/blob/c612628045822f909020f7eb6784c79700813eda/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L410
+            "context_length": "max_position_embeddings",
+            "vocab_size": "vocab_size",
+            "width": "hidden_size",
+            "heads": "num_attention_heads",
+            "layers": "num_hidden_layers",
+            "layer_attr": "layers",
+            "token_embeddings_attr": "embed_in"
+        },
+        "pooler": "weighted_mean_pooler",
+    },
 }