Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+vitl14_attention.png filter=lfs diff=lfs merge=lfs -text
+vitl14_patchnorms.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,163 @@
+---
+library_name: transformers
+license: mit
+pipeline_tag: image-feature-extraction
+tags:
+- clip
+---
+
+# OpenCLIP ViT-L/14 with Test-Time Register
+
+Register tokens in ViTs were introduced as learnable tokens in [Vision Transformers Need Registers](https://arxiv.org/abs/2309.16588) to mitigate artifacts in intermediate feature maps.
+In [Vision Transformers Don't Need *Trained* Registers](https://arxiv.org/abs/2506.08010), we introduced a training-free method to create registers. These *test-time registers* serve a similar purpose
+as the original trained registers, but can be added post-hoc to any ViT to mitigate artifacts, enhance model interpretability, and modestly improve downstream performance in tasks such as segmentation, depth estimation, etc. 
+
+## Model description
+
+The base model is [OpenCLIP-ViT-L-14-laion2B-s32B-b82K](https://huggingface.co/laion/CLIP-ViT-L-14-laion2B-s32B-b82K). With test-time registers, the model's internal representations
+are cleaner (see below). Using the environment from [here](https://github.com/nickjiang2378/test-time-registers/blob/main/environment.yml) and evaluating using bfloat16 leads to IN-1k zeroshot performance of 76.4 for both the original model and the variant with test-time registers.
+This model is intended to be used with this [repo](https://github.com/nickjiang2378/test-time-registers). Use transformers==4.45.1. The model can also be used for fine-tuning or other downstream tasks.
+
+<img src="https://huggingface.co/amildravid4292/clip-vitl14-test-time-registers/resolve/main/vitl14_attention.png" alt="drawing" width="600"/>
+<img src="https://huggingface.co/amildravid4292/clip-vitl14-test-time-registers/resolve/main/vitl14_patchnorms.png" alt="drawing" width="600"/>
+
+
+
+
+## Quick Start
+
+```python
+from transformers import AutoModel
+from PIL import Image
+import torch
+
+# Load the complete model with all components
+model = AutoModel.from_pretrained(
+    "amildravid4292/clip-vitl14-test-time-registers", 
+    trust_remote_code=True
+)
+
+# Check what was loaded
+print(f"Register tokens: {model.num_register_tokens}")
+print(f"Neuron dict: {model.neuron_dict}")
+print(f"Tokenizer available: {model.tokenizer is not None}")
+print(f"Preprocessor available: {model.preprocessor is not None}")
+print(f"Zero-shot classifier available: {model.zeroshot_classifier is not None}")
+```
+
+## Usage Examples
+
+
+
+### Image Processing
+```python
+from PIL import Image
+
+# Load and preprocess image
+image = Image.open("your_image.jpg")
+image_tensor = model.preprocess_image(image).unsqueeze(0)
+
+image_features = model.encode_image(
+    image_tensor
+)
+
+# to run inference with the original model without test-time registers
+image_features = model.encode_image(
+    image_tensor,
+    neuron_dict=None,
+    num_register_tokens=0
+)
+
+```
+
+### Text Processing
+```python
+# Tokenize text
+text = ["a photo of a cat", "a photo of a dog"]
+text_tokens = model.tokenize(text)
+
+# Encode text
+text_features = model.encode_text(text_tokens)
+```
+
+
+
+### Complete Pipeline
+```python
+
+# load model
+model = AutoModel.from_pretrained('amildravid4292/clip-vitl14-test-time-registers', trust_remote_code=True)
+model = model.to(device).bfloat16()
+classifier = model.zeroshot_classifier.to(device).bfloat16()
+
+# load data
+imagenet_dataset = ImageNet(root='/datasets/ilsvrc/current', split='val', transform=model.preprocessor)
+ground_truth_labels = [imagenet_dataset.targets[i] for i in range(len(imagenet_dataset))]
+loader = torch.utils.data.DataLoader(imagenet_dataset, batch_size=100, num_workers=4, pin_memory=True, shuffle=False)
+
+# run zero-shot classification
+with torch.no_grad():
+    correct = [0, 0]
+    for i, (images, target) in enumerate(tqdm(loader)):
+        images = images.to(device).bfloat16()
+        
+        target = target.to(device).bfloat16()
+    
+        
+        # predict
+        image_features = model.encode_image(images) 
+        
+        image_features /= image_features.norm(dim=-1, keepdim=True)
+        logits = 100. * image_features @ classifier
+
+        pred = logits.argmax(dim=-1)
+        correct[0] += (pred == target).sum().item()
+        correct[1] += target.size(0)
+        
+       
+        
+print(correct[0]/correct[1])
+```
+
+## Advanced Usage
+
+### Custom Neuron Modifications
+```python
+# Override the saved neuron configuration
+custom_neuron_dict = {0: [10, 20, 30]}  # Modify neurons 10,20,30 in layer 0
+
+image_features = model.encode_image(
+    image_tensor,
+    num_register_tokens=4,
+    neuron_dict=custom_neuron_dict
+)
+```
+
+### Different Register Token Counts
+```python
+# Use different number of register tokens
+image_features = model.encode_image(
+    image_tensor,
+    num_register_tokens=8  # Override the default
+)
+```
+
+## Model Details
+
+- **Base Architecture**: ViT-L/14
+- **Training Data**: LAION-2B subset
+
+
+### BibTeX entry and citation info
+
+```bibtex
+@misc{jiang2025visiontransformersdontneed,
+      title={Vision Transformers Don't Need Trained Registers}, 
+      author={Nick Jiang and Amil Dravid and Alexei Efros and Yossi Gandelsman},
+      year={2025},
+      eprint={2506.08010},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2506.08010}, 
+}
+```

ViT-L-14-336_register.json

@@ -0,0 +1,116 @@
+{
+    "9": [
+        815,
+        4078,
+        3618,
+        2693,
+        3973,
+        1744,
+        1983,
+        1157,
+        1309,
+        1335,
+        2607,
+        2396,
+        3049,
+        1610,
+        2621,
+        2867,
+        2012,
+        1924,
+        2394,
+        3097,
+        3125,
+        3959,
+        3210,
+        2855,
+        3609,
+        526,
+        3362,
+        3395,
+        2626,
+        503,
+        2941,
+        3696,
+        1823,
+        2000,
+        129,
+        3667,
+        1372,
+        147,
+        1150,
+        852,
+        3222
+    ],
+    "8": [
+        745,
+        3249,
+        2585,
+        1537,
+        200,
+        1603,
+        1851,
+        3523,
+        3697,
+        3137,
+        2563,
+        2293,
+        730,
+        906,
+        1528,
+        3348,
+        2438,
+        1564,
+        1540,
+        3238,
+        3606
+    ],
+    "10": [
+        357,
+        1654,
+        3940,
+        2319,
+        2560,
+        2559,
+        4009,
+        3029,
+        951,
+        1903,
+        738,
+        1602,
+        1807,
+        2018,
+        1281,
+        267,
+        3539,
+        1015,
+        496,
+        693,
+        2278,
+        7,
+        856,
+        2785,
+        2690,
+        1367
+    ],
+    "7": [
+        3228,
+        2550,
+        2977,
+        3716,
+        2467
+    ],
+    "0": [
+        2890,
+        1779,
+        3761
+    ],
+    "6": [
+        1042,
+        2315,
+        1674
+    ],
+    "3": [
+        410
+    ]
+}

__init__.py

@@ -0,0 +1,8 @@
+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 pretrained import list_pretrained, list_pretrained_models_by_tag, list_pretrained_tags_by_model, \
+    get_pretrained_url, download_pretrained_from_url, is_pretrained_cfg, get_pretrained_cfg, download_pretrained
+from tokenizer import SimpleTokenizer, tokenize, decode
+from transform import image_transform, AugmentationCfg
+from openai_templates import OPENAI_IMAGENET_TEMPLATES

config.json

@@ -0,0 +1,29 @@
+{
+  "model_type": "custom_clip_with_registers",
+
+  "processor_class": "CLIPProcessor",
+  "tokenizer_class": "CLIPTokenizerFast",
+
+  "architectures": ["CustomCLIPModel"],
+  "auto_map": {
+    "AutoConfig": "modeling_custom_clip.CustomCLIPConfig", 
+    "AutoModel": "modeling_custom_clip.CustomCLIPModel"
+  },
+  "vision_config": {
+    "hidden_size": 1024,
+    "num_hidden_layers": 24,
+    "num_attention_heads": 16,
+    "image_size": 336,
+    "patch_size": 14
+  },
+  "text_config": {
+    "vocab_size": 49408,
+    "hidden_size": 768,
+    "num_hidden_layers": 12,
+    "max_position_embeddings": 77
+  },
+  "neuron_dict": {},
+  "projection_dim": 768,      
+  "torch_dtype": "float32",
+  "transformers_version": "4.21.0"
+}

config_TTR_bak.json

@@ -0,0 +1,145 @@
+{
+  "model_type": "custom_clip_with_registers",
+
+  "processor_class": "CLIPProcessor",
+  "tokenizer_class": "CLIPTokenizerFast",
+
+  "architectures": ["CustomCLIPModel"],
+  "auto_map": {
+    "AutoConfig": "modeling_custom_clip.CustomCLIPConfig", 
+    "AutoModel": "modeling_custom_clip.CustomCLIPModel"
+  },
+  "vision_config": {
+    "hidden_size": 1024,
+    "num_hidden_layers": 24,
+    "num_attention_heads": 16,
+    "image_size": 336,
+    "patch_size": 14
+  },
+  "text_config": {
+    "vocab_size": 49408,
+    "hidden_size": 768,
+    "num_hidden_layers": 12,
+    "max_position_embeddings": 77
+  },
+  "num_register_tokens": 1,
+  "neuron_dict": {
+    "9": [
+        815,
+        4078,
+        3618,
+        2693,
+        3973,
+        1744,
+        1983,
+        1157,
+        1309,
+        1335,
+        2607,
+        2396,
+        3049,
+        1610,
+        2621,
+        2867,
+        2012,
+        1924,
+        2394,
+        3097,
+        3125,
+        3959,
+        3210,
+        2855,
+        3609,
+        526,
+        3362,
+        3395,
+        2626,
+        503,
+        2941,
+        3696,
+        1823,
+        2000,
+        129,
+        3667,
+        1372,
+        147,
+        1150,
+        852,
+        3222
+    ],
+    "8": [
+        745,
+        3249,
+        2585,
+        1537,
+        200,
+        1603,
+        1851,
+        3523,
+        3697,
+        3137,
+        2563,
+        2293,
+        730,
+        906,
+        1528,
+        3348,
+        2438,
+        1564,
+        1540,
+        3238,
+        3606
+    ],
+    "10": [
+        357,
+        1654,
+        3940,
+        2319,
+        2560,
+        2559,
+        4009,
+        3029,
+        951,
+        1903,
+        738,
+        1602,
+        1807,
+        2018,
+        1281,
+        267,
+        3539,
+        1015,
+        496,
+        693,
+        2278,
+        7,
+        856,
+        2785,
+        2690,
+        1367
+    ],
+    "7": [
+        3228,
+        2550,
+        2977,
+        3716,
+        2467
+    ],
+    "0": [
+        2890,
+        1779,
+        3761
+    ],
+    "6": [
+        1042,
+        2315,
+        1674
+    ],
+    "3": [
+        410
+    ]
+},
+  "projection_dim": 768,      
+  "torch_dtype": "float32",
+  "transformers_version": "4.21.0"
+}

config_bak.json

@@ -0,0 +1,110 @@
+{
+  "model_type": "custom_clip_with_registers",
+  "architectures": ["CustomCLIPModel"],
+  "auto_map": {
+    "AutoConfig": "modeling_custom_clip.CustomCLIPConfig", 
+    "AutoModel": "modeling_custom_clip.CustomCLIPModel"
+  },
+  "vision_config": {
+    "hidden_size": 1024,
+    "num_hidden_layers": 24,
+    "num_attention_heads": 16,
+    "image_size": 336,
+    "patch_size": 14
+  },
+  "text_config": {
+    "vocab_size": 49408,
+    "hidden_size": 768,
+    "num_hidden_layers": 12,
+    "max_position_embeddings": 77
+  },
+  "num_register_tokens": 1,
+  "neuron_dict": {"10": [2924,
+  2520,
+  2936,
+  675,
+  517,
+  1610,
+  88,
+  1950,
+  3098,
+  4082,
+  1237,
+  857,
+  3020,
+  1321,
+  1128,
+  3561,
+  4091,
+  69,
+  3378,
+  2304,
+  977,
+  1762,
+  3598,
+  371,
+  1097],
+ "9": [1253,
+  3658,
+  1827,
+  2600,
+  4000,
+  711,
+  2726,
+  615,
+  2654,
+  831,
+  1,
+  1387,
+  2178,
+  1967,
+  2413,
+  901,
+  481,
+  1514,
+  292,
+  692,
+  3094,
+  3470,
+  932,
+  2129],
+ "8": [3189,
+  1491,
+  2159,
+  1196,
+  1913,
+  1340,
+  2515,
+  2163,
+  955,
+  1496,
+  1891,
+  1410,
+  3725,
+  632,
+  188,
+  726,
+  1592,
+  1017,
+  1267,
+  995,
+  3465,
+  3510,
+  1494,
+  3467,
+  1896,
+  2779,
+  2309,
+  3389,
+  3682,
+  1968,
+  2904],
+ "7": [2226, 2565],
+ "6": [1450, 1551, 1024],
+ "5": [151, 1282],
+ "4": [2207],
+ "3": [2298, 2841]},
+  "projection_dim": 768,      
+  "torch_dtype": "float32",
+  "transformers_version": "4.21.0"
+}

constants.py

@@ -0,0 +1,2 @@
+OPENAI_DATASET_MEAN = (0.48145466, 0.4578275, 0.40821073)
+OPENAI_DATASET_STD = (0.26862954, 0.26130258, 0.27577711)

factory.py

@@ -0,0 +1,390 @@
+import json
+import logging
+import os
+import pathlib
+import re
+from copy import deepcopy
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+
+from constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from model import CLIP, convert_to_custom_text_state_dict,\
+    resize_pos_embed, get_cast_dtype
+from openai_models import load_openai_model
+from pretrained import is_pretrained_cfg, get_pretrained_cfg, download_pretrained,\
+    list_pretrained_tags_by_model, download_pretrained_from_hf
+from transform import image_transform, AugmentationCfg
+from tokenizer import HFTokenizer, tokenize
+
+
+HF_HUB_PREFIX = 'hf-hub:'
+_MODEL_CONFIG_PATHS = [Path(__file__).parent / f"model_configs/"]
+_MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
+
+
+def _natural_key(string_):
+    return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
+
+
+def _rescan_model_configs():
+    global _MODEL_CONFIGS
+
+    config_ext = ('.json',)
+    config_files = []
+    for config_path in _MODEL_CONFIG_PATHS:
+        if config_path.is_file() and config_path.suffix in config_ext:
+            config_files.append(config_path)
+        elif config_path.is_dir():
+            for ext in config_ext:
+                config_files.extend(config_path.glob(f'*{ext}'))
+
+    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')):
+                _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
+
+
+def list_models():
+    """ enumerate available model architectures based on config files """
+    return list(_MODEL_CONFIGS.keys())
+
+
+def add_model_config(path):
+    """ add model config path or file and update registry """
+    if not isinstance(path, Path):
+        path = Path(path)
+    _MODEL_CONFIG_PATHS.append(path)
+    _rescan_model_configs()
+
+
+def get_model_config(model_name):
+    if model_name in _MODEL_CONFIGS:
+        return deepcopy(_MODEL_CONFIGS[model_name])
+    else:
+        return None
+
+
+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
+    return tokenizer
+
+
+def load_state_dict(checkpoint_path: str, map_location='cpu'):
+    checkpoint = torch.load(checkpoint_path, map_location=map_location)
+    if isinstance(checkpoint, dict) and 'state_dict' in checkpoint:
+        state_dict = checkpoint['state_dict']
+    else:
+        state_dict = checkpoint
+    if next(iter(state_dict.items()))[0].startswith('module'):
+        state_dict = {k[7:]: v for k, v in state_dict.items()}
+    return state_dict
+
+
+def load_checkpoint(model, checkpoint_path, strict=False):
+    state_dict = load_state_dict(checkpoint_path)
+    # detect old format and make compatible with new format
+    if 'positional_embedding' in state_dict and not hasattr(model, 'positional_embedding'):
+        state_dict = convert_to_custom_text_state_dict(state_dict)
+    resize_pos_embed(state_dict, model)
+    incompatible_keys = model.load_state_dict(state_dict, strict=strict)
+
+    model.num_register_tokens=state_dict["num_register_tokens"]
+    model.neuron_dict=state_dict["neuron_dict"]
+    model.visual.num_register_tokens=state_dict["num_register_tokens"]
+    model.visual.neuron_dict=state_dict["neuron_dict"]
+    
+    return incompatible_keys
+
+
+def create_model(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        cache_dir: Optional[str] = None,
+        output_dict: Optional[bool] = None,
+        require_pretrained: bool = False,
+):
+    has_hf_hub_prefix = model_name.startswith(HF_HUB_PREFIX)
+    if has_hf_hub_prefix:
+        model_id = model_name[len(HF_HUB_PREFIX):]
+        checkpoint_path = download_pretrained_from_hf(model_id, cache_dir=cache_dir)
+        config_path = download_pretrained_from_hf(model_id, filename='open_clip_config.json', cache_dir=cache_dir)
+        
+        with open(config_path, 'r', encoding='utf-8') as f:
+            config = json.load(f)
+        pretrained_cfg = config['preprocess_cfg']
+        model_cfg = config['model_cfg']
+    else:
+        model_name = model_name.replace('/', '-')  # for callers using old naming with / in ViT names
+        checkpoint_path = None
+        pretrained_cfg = {}
+        model_cfg = None
+        
+        
+
+    if isinstance(device, str):
+        device = torch.device(device)
+
+    if pretrained and pretrained.lower() == 'openai':
+        logging.info(f'Loading pretrained {model_name} from OpenAI.')
+        model = load_openai_model(
+            model_name,
+            precision=precision,
+            device=device,
+            cache_dir=cache_dir,
+            quick_gelu=force_quick_gelu,
+        )
+    else:
+        model_cfg = model_cfg or get_model_config(model_name)
+        
+        if model_cfg is not None:
+            logging.info(f'Loaded {model_name} model config.')
+        else:
+            logging.error(f'Model config for {model_name} not found; available models {list_models()}.')
+            raise RuntimeError(f'Model config for {model_name} not found.')
+
+        if force_patch_dropout is not None:
+            # override the default patch dropout value
+            model_cfg["vision_cfg"]["patch_dropout"] = force_patch_dropout
+
+        if force_image_size is not None:
+            # override model config's image size
+            model_cfg["vision_cfg"]["image_size"] = force_image_size
+
+        is_timm_model = 'timm_model_name' in model_cfg.get('vision_cfg', {})
+        if pretrained_image:
+            if is_timm_model:
+                # pretrained weight loading for timm models set via vision_cfg
+                model_cfg['vision_cfg']['timm_model_pretrained'] = True
+            else:
+                assert False, 'pretrained image towers currently only supported for timm models'
+
+        # cast_dtype set for fp16 and bf16 (manual mixed-precision), not set for 'amp' or 'pure' modes
+        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:
+                raise ValueError('Coca is not implemented')
+                model = CoCa(**model_cfg, cast_dtype=cast_dtype)
+            else:
+                raise ValueError('CustomTextCLIP is not implemented')
+                model = CustomTextCLIP(**model_cfg, cast_dtype=cast_dtype)
+        else:
+            model = CLIP(**model_cfg, cast_dtype=cast_dtype)
+
+        if precision in ("fp16", "bf16"):
+            dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
+            # manual mixed precision that matches original OpenAI behaviour
+            if is_timm_model:
+                # FIXME this is a bit janky, create timm based model in low-precision and
+                # then cast only LayerNormFp32 instances back to float32 so they don't break.
+                # Why? The convert_weights_to_lp fn only works with native models.
+                model.to(device=device, dtype=dtype)
+                from transformer import LayerNormFp32
+                def _convert_ln(m):
+                    if isinstance(m, LayerNormFp32):
+                        m.weight.data = m.weight.data.to(torch.float32)
+                        m.bias.data = m.bias.data.to(torch.float32)
+                model.apply(_convert_ln)
+            else:
+                model.to(device=device)
+                convert_weights_to_lp(model, dtype=dtype)
+        elif precision in ("pure_fp16", "pure_bf16"):
+            dtype = torch.float16 if 'fp16' in precision else torch.bfloat16
+            model.to(device=device, dtype=dtype)
+        else:
+            model.to(device=device)
+
+        pretrained_loaded = False
+        if pretrained:
+            checkpoint_path = ''
+            pretrained_cfg = get_pretrained_cfg(model_name, pretrained)
+            if pretrained_cfg:
+                checkpoint_path = download_pretrained(pretrained_cfg, cache_dir=cache_dir)
+            elif os.path.exists(pretrained):
+                checkpoint_path = pretrained
+
+            if checkpoint_path:
+                logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+                load_checkpoint(model, checkpoint_path)
+            else:
+                error_str = (
+                    f'Pretrained weights ({pretrained}) not found for model {model_name}.'
+                    f'Available pretrained tags ({list_pretrained_tags_by_model(model_name)}.')
+                logging.warning(error_str)
+                raise RuntimeError(error_str)
+            pretrained_loaded = True
+        elif has_hf_hub_prefix:
+            logging.info(f'Loading pretrained {model_name} weights ({pretrained}).')
+            load_checkpoint(model, checkpoint_path)
+            pretrained_loaded = True
+
+        if require_pretrained and not pretrained_loaded:
+            # callers of create_model_from_pretrained always expect pretrained weights
+            raise RuntimeError(
+                f'Pretrained weights were required for (model: {model_name}, pretrained: {pretrained}) but not loaded.')
+
+        # 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
+
+    if output_dict and hasattr(model, "output_dict"):
+        model.output_dict = True
+
+    if jit:
+        model = torch.jit.script(model)
+
+    return model
+
+
+def create_loss(args):
+    if args.distill:
+        return DistillClipLoss(
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+            use_horovod=args.horovod,
+        )
+    elif "coca" in args.model.lower():
+        return CoCaLoss(
+            caption_loss_weight=args.coca_caption_loss_weight,
+            clip_loss_weight=args.coca_contrastive_loss_weight,
+            local_loss=args.local_loss,
+            gather_with_grad=args.gather_with_grad,
+            cache_labels=True,
+            rank=args.rank,
+            world_size=args.world_size,
+            use_horovod=args.horovod,
+        )
+    return ClipLoss(
+        local_loss=args.local_loss,
+        gather_with_grad=args.gather_with_grad,
+        cache_labels=True,
+        rank=args.rank,
+        world_size=args.world_size,
+        use_horovod=args.horovod,
+    )
+
+
+def create_model_and_transforms(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_patch_dropout: Optional[float] = None,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        pretrained_image: bool = False,
+        pretrained_hf: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
+        cache_dir: Optional[str] = None,
+        output_dict: Optional[bool] = None,
+):
+    
+    
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_patch_dropout=force_patch_dropout,
+        force_image_size=force_image_size,
+        pretrained_image=pretrained_image,
+        pretrained_hf=pretrained_hf,
+        cache_dir=cache_dir,
+        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,
+    )
+
+    return model, preprocess_train, preprocess_val
+
+
+def create_model_from_pretrained(
+        model_name: str,
+        pretrained: Optional[str] = None,
+        precision: str = 'fp32',
+        device: Union[str, torch.device] = 'cpu',
+        jit: bool = False,
+        force_quick_gelu: bool = False,
+        force_custom_text: bool = False,
+        force_image_size: Optional[Union[int, Tuple[int, int]]] = None,
+        return_transform: bool = True,
+        image_mean: Optional[Tuple[float, ...]] = None,
+        image_std: Optional[Tuple[float, ...]] = None,
+        cache_dir: Optional[str] = None,
+):
+    model = create_model(
+        model_name,
+        pretrained,
+        precision=precision,
+        device=device,
+        jit=jit,
+        force_quick_gelu=force_quick_gelu,
+        force_custom_text=force_custom_text,
+        force_image_size=force_image_size,
+        cache_dir=cache_dir,
+        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,
+    )
+
+    return model, preprocess

imagenet_classes.py

@@ -0,0 +1,85 @@
+OPENAI_IMAGENET_TEMPLATES = (
+    lambda c: f'a bad photo of a {c}.',
+    lambda c: f'a photo of many {c}.',
+    lambda c: f'a sculpture of a {c}.',
+    lambda c: f'a photo of the hard to see {c}.',
+    lambda c: f'a low resolution photo of the {c}.',
+    lambda c: f'a rendering of a {c}.',
+    lambda c: f'graffiti of a {c}.',
+    lambda c: f'a bad photo of the {c}.',
+    lambda c: f'a cropped photo of the {c}.',
+    lambda c: f'a tattoo of a {c}.',
+    lambda c: f'the embroidered {c}.',
+    lambda c: f'a photo of a hard to see {c}.',
+    lambda c: f'a bright photo of a {c}.',
+    lambda c: f'a photo of a clean {c}.',
+    lambda c: f'a photo of a dirty {c}.',
+    lambda c: f'a dark photo of the {c}.',
+    lambda c: f'a drawing of a {c}.',
+    lambda c: f'a photo of my {c}.',
+    lambda c: f'the plastic {c}.',
+    lambda c: f'a photo of the cool {c}.',
+    lambda c: f'a close-up photo of a {c}.',
+    lambda c: f'a black and white photo of the {c}.',
+    lambda c: f'a painting of the {c}.',
+    lambda c: f'a painting of a {c}.',
+    lambda c: f'a pixelated photo of the {c}.',
+    lambda c: f'a sculpture of the {c}.',
+    lambda c: f'a bright photo of the {c}.',
+    lambda c: f'a cropped photo of a {c}.',
+    lambda c: f'a plastic {c}.',
+    lambda c: f'a photo of the dirty {c}.',
+    lambda c: f'a jpeg corrupted photo of a {c}.',
+    lambda c: f'a blurry photo of the {c}.',
+    lambda c: f'a photo of the {c}.',
+    lambda c: f'a good photo of the {c}.',
+    lambda c: f'a rendering of the {c}.',
+    lambda c: f'a {c} in a video game.',
+    lambda c: f'a photo of one {c}.',
+    lambda c: f'a doodle of a {c}.',
+    lambda c: f'a close-up photo of the {c}.',
+    lambda c: f'a photo of a {c}.',
+    lambda c: f'the origami {c}.',
+    lambda c: f'the {c} in a video game.',
+    lambda c: f'a sketch of a {c}.',
+    lambda c: f'a doodle of the {c}.',
+    lambda c: f'a origami {c}.',
+    lambda c: f'a low resolution photo of a {c}.',
+    lambda c: f'the toy {c}.',
+    lambda c: f'a rendition of the {c}.',
+    lambda c: f'a photo of the clean {c}.',
+    lambda c: f'a photo of a large {c}.',
+    lambda c: f'a rendition of a {c}.',
+    lambda c: f'a photo of a nice {c}.',
+    lambda c: f'a photo of a weird {c}.',
+    lambda c: f'a blurry photo of a {c}.',
+    lambda c: f'a cartoon {c}.',
+    lambda c: f'art of a {c}.',
+    lambda c: f'a sketch of the {c}.',
+    lambda c: f'a embroidered {c}.',
+    lambda c: f'a pixelated photo of a {c}.',
+    lambda c: f'itap of the {c}.',
+    lambda c: f'a jpeg corrupted photo of the {c}.',
+    lambda c: f'a good photo of a {c}.',
+    lambda c: f'a plushie {c}.',
+    lambda c: f'a photo of the nice {c}.',
+    lambda c: f'a photo of the small {c}.',
+    lambda c: f'a photo of the weird {c}.',
+    lambda c: f'the cartoon {c}.',
+    lambda c: f'art of the {c}.',
+    lambda c: f'a drawing of the {c}.',
+    lambda c: f'a photo of the large {c}.',
+    lambda c: f'a black and white photo of a {c}.',
+    lambda c: f'the plushie {c}.',
+    lambda c: f'a dark photo of a {c}.',
+    lambda c: f'itap of a {c}.',
+    lambda c: f'graffiti of the {c}.',
+    lambda c: f'a toy {c}.',
+    lambda c: f'itap of my {c}.',
+    lambda c: f'a photo of a cool {c}.',
+    lambda c: f'a photo of a small {c}.',
+    lambda c: f'a tattoo of the {c}.',
+)
+
+
+IMAGENET_CLASSNAMES = ["tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark", "electric ray", "stingray", "rooster", "hen", "ostrich", "brambling", "goldfinch", "house finch", "junco", "indigo bunting", "American robin", "bulbul", "jay", "magpie", "chickadee", "American dipper", "kite (bird of prey)", "bald eagle", "vulture", "great grey owl", "fire salamander", "smooth newt", "newt", "spotted salamander", "axolotl", "American bullfrog", "tree frog", "tailed frog", "loggerhead sea turtle", "leatherback sea turtle", "mud turtle", "terrapin", "box turtle", "banded gecko", "green iguana", "Carolina anole", "desert grassland whiptail lizard", "agama", "frilled-necked lizard", "alligator lizard", "Gila monster", "European green lizard", "chameleon", "Komodo dragon", "Nile crocodile", "American alligator", "triceratops", "worm snake", "ring-necked snake", "eastern hog-nosed snake", "smooth green snake", "kingsnake", "garter snake", "water snake", "vine snake", "night snake", "boa constrictor", "African rock python", "Indian cobra", "green mamba", "sea snake", "Saharan horned viper", "eastern diamondback rattlesnake", "sidewinder rattlesnake", "trilobite", "harvestman", "scorpion", "yellow garden spider", "barn spider", "European garden spider", "southern black widow", "tarantula", "wolf spider", "tick", "centipede", "black grouse", "ptarmigan", "ruffed grouse", "prairie grouse", "peafowl", "quail", "partridge", "african grey parrot", "macaw", "sulphur-crested cockatoo", "lorikeet", "coucal", "bee eater", "hornbill", "hummingbird", "jacamar", "toucan", "duck", "red-breasted merganser", "goose", "black swan", "tusker", "echidna", "platypus", "wallaby", "koala", "wombat", "jellyfish", "sea anemone", "brain coral", "flatworm", "nematode", "conch", "snail", "slug", "sea slug", "chiton", "chambered nautilus", "Dungeness crab", "rock crab", "fiddler crab", "red king crab", "American lobster", "spiny lobster", "crayfish", "hermit crab", "isopod", "white stork", "black stork", "spoonbill", "flamingo", "little blue heron", "great egret", "bittern bird", "crane bird", "limpkin", "common gallinule", "American coot", "bustard", "ruddy turnstone", "dunlin", "common redshank", "dowitcher", "oystercatcher", "pelican", "king penguin", "albatross", "grey whale", "killer whale", "dugong", "sea lion", "Chihuahua", "Japanese Chin", "Maltese", "Pekingese", "Shih Tzu", "King Charles Spaniel", "Papillon", "toy terrier", "Rhodesian Ridgeback", "Afghan Hound", "Basset Hound", "Beagle", "Bloodhound", "Bluetick Coonhound", "Black and Tan Coonhound", "Treeing Walker Coonhound", "English foxhound", "Redbone Coonhound", "borzoi", "Irish Wolfhound", "Italian Greyhound", "Whippet", "Ibizan Hound", "Norwegian Elkhound", "Otterhound", "Saluki", "Scottish Deerhound", "Weimaraner", "Staffordshire Bull Terrier", "American Staffordshire Terrier", "Bedlington Terrier", "Border Terrier", "Kerry Blue Terrier", "Irish Terrier", "Norfolk Terrier", "Norwich Terrier", "Yorkshire Terrier", "Wire Fox Terrier", "Lakeland Terrier", "Sealyham Terrier", "Airedale Terrier", "Cairn Terrier", "Australian Terrier", "Dandie Dinmont Terrier", "Boston Terrier", "Miniature Schnauzer", "Giant Schnauzer", "Standard Schnauzer", "Scottish Terrier", "Tibetan Terrier", "Australian Silky Terrier", "Soft-coated Wheaten Terrier", "West Highland White Terrier", "Lhasa Apso", "Flat-Coated Retriever", "Curly-coated Retriever", "Golden Retriever", "Labrador Retriever", "Chesapeake Bay Retriever", "German Shorthaired Pointer", "Vizsla", "English Setter", "Irish Setter", "Gordon Setter", "Brittany dog", "Clumber Spaniel", "English Springer Spaniel", "Welsh Springer Spaniel", "Cocker Spaniel", "Sussex Spaniel", "Irish Water Spaniel", "Kuvasz", "Schipperke", "Groenendael dog", "Malinois", "Briard", "Australian Kelpie", "Komondor", "Old English Sheepdog", "Shetland Sheepdog", "collie", "Border Collie", "Bouvier des Flandres dog", "Rottweiler", "German Shepherd Dog", "Dobermann", "Miniature Pinscher", "Greater Swiss Mountain Dog", "Bernese Mountain Dog", "Appenzeller Sennenhund", "Entlebucher Sennenhund", "Boxer", "Bullmastiff", "Tibetan Mastiff", "French Bulldog", "Great Dane", "St. Bernard", "husky", "Alaskan Malamute", "Siberian Husky", "Dalmatian", "Affenpinscher", "Basenji", "pug", "Leonberger", "Newfoundland dog", "Great Pyrenees dog", "Samoyed", "Pomeranian", "Chow Chow", "Keeshond", "brussels griffon", "Pembroke Welsh Corgi", "Cardigan Welsh Corgi", "Toy Poodle", "Miniature Poodle", "Standard Poodle", "Mexican hairless dog (xoloitzcuintli)", "grey wolf", "Alaskan tundra wolf", "red wolf or maned wolf", "coyote", "dingo", "dhole", "African wild dog", "hyena", "red fox", "kit fox", "Arctic fox", "grey fox", "tabby cat", "tiger cat", "Persian cat", "Siamese cat", "Egyptian Mau", "cougar", "lynx", "leopard", "snow leopard", "jaguar", "lion", "tiger", "cheetah", "brown bear", "American black bear", "polar bear", "sloth bear", "mongoose", "meerkat", "tiger beetle", "ladybug", "ground beetle", "longhorn beetle", "leaf beetle", "dung beetle", "rhinoceros beetle", "weevil", "fly", "bee", "ant", "grasshopper", "cricket insect", "stick insect", "cockroach", "praying mantis", "cicada", "leafhopper", "lacewing", "dragonfly", "damselfly", "red admiral butterfly", "ringlet butterfly", "monarch butterfly", "small white butterfly", "sulphur butterfly", "gossamer-winged butterfly", "starfish", "sea urchin", "sea cucumber", "cottontail rabbit", "hare", "Angora rabbit", "hamster", "porcupine", "fox squirrel", "marmot", "beaver", "guinea pig", "common sorrel horse", "zebra", "pig", "wild boar", "warthog", "hippopotamus", "ox", "water buffalo", "bison", "ram (adult male sheep)", "bighorn sheep", "Alpine ibex", "hartebeest", "impala (antelope)", "gazelle", "arabian camel", "llama", "weasel", "mink", "European polecat", "black-footed ferret", "otter", "skunk", "badger", "armadillo", "three-toed sloth", "orangutan", "gorilla", "chimpanzee", "gibbon", "siamang", "guenon", "patas monkey", "baboon", "macaque", "langur", "black-and-white colobus", "proboscis monkey", "marmoset", "white-headed capuchin", "howler monkey", "titi monkey", "Geoffroy's spider monkey", "common squirrel monkey", "ring-tailed lemur", "indri", "Asian elephant", "African bush elephant", "red panda", "giant panda", "snoek fish", "eel", "silver salmon", "rock beauty fish", "clownfish", "sturgeon", "gar fish", "lionfish", "pufferfish", "abacus", "abaya", "academic gown", "accordion", "acoustic guitar", "aircraft carrier", "airliner", "airship", "altar", "ambulance", "amphibious vehicle", "analog clock", "apiary", "apron", "trash can", "assault rifle", "backpack", "bakery", "balance beam", "balloon", "ballpoint pen", "Band-Aid", "banjo", "baluster / handrail", "barbell", "barber chair", "barbershop", "barn", "barometer", "barrel", "wheelbarrow", "baseball", "basketball", "bassinet", "bassoon", "swimming cap", "bath towel", "bathtub", "station wagon", "lighthouse", "beaker", "military hat (bearskin or shako)", "beer bottle", "beer glass", "bell tower", "baby bib", "tandem bicycle", "bikini", "ring binder", "binoculars", "birdhouse", "boathouse", "bobsleigh", "bolo tie", "poke bonnet", "bookcase", "bookstore", "bottle cap", "hunting bow", "bow tie", "brass memorial plaque", "bra", "breakwater", "breastplate", "broom", "bucket", "buckle", "bulletproof vest", "high-speed train", "butcher shop", "taxicab", "cauldron", "candle", "cannon", "canoe", "can opener", "cardigan", "car mirror", "carousel", "tool kit", "cardboard box / carton", "car wheel", "automated teller machine", "cassette", "cassette player", "castle", "catamaran", "CD player", "cello", "mobile phone", "chain", "chain-link fence", "chain mail", "chainsaw", "storage chest", "chiffonier", "bell or wind chime", "china cabinet", "Christmas stocking", "church", "movie theater", "cleaver", "cliff dwelling", "cloak", "clogs", "cocktail shaker", "coffee mug", "coffeemaker", "spiral or coil", "combination lock", "computer keyboard", "candy store", "container ship", "convertible", "corkscrew", "cornet", "cowboy boot", "cowboy hat", "cradle", "construction crane", "crash helmet", "crate", "infant bed", "Crock Pot", "croquet ball", "crutch", "cuirass", "dam", "desk", "desktop computer", "rotary dial telephone", "diaper", "digital clock", "digital watch", "dining table", "dishcloth", "dishwasher", "disc brake", "dock", "dog sled", "dome", "doormat", "drilling rig", "drum", "drumstick", "dumbbell", "Dutch oven", "electric fan", "electric guitar", "electric locomotive", "entertainment center", "envelope", "espresso machine", "face powder", "feather boa", "filing cabinet", "fireboat", "fire truck", "fire screen", "flagpole", "flute", "folding chair", "football helmet", "forklift", "fountain", "fountain pen", "four-poster bed", "freight car", "French horn", "frying pan", "fur coat", "garbage truck", "gas mask or respirator", "gas pump", "goblet", "go-kart", "golf ball", "golf cart", "gondola", "gong", "gown", "grand piano", "greenhouse", "radiator grille", "grocery store", "guillotine", "hair clip", "hair spray", "half-track", "hammer", "hamper", "hair dryer", "hand-held computer", "handkerchief", "hard disk drive", "harmonica", "harp", "combine harvester", "hatchet", "holster", "home theater", "honeycomb", "hook", "hoop skirt", "gymnastic horizontal bar", "horse-drawn vehicle", "hourglass", "iPod", "clothes iron", "carved pumpkin", "jeans", "jeep", "T-shirt", "jigsaw puzzle", "rickshaw", "joystick", "kimono", "knee pad", "knot", "lab coat", "ladle", "lampshade", "laptop computer", "lawn mower", "lens cap", "letter opener", "library", "lifeboat", "lighter", "limousine", "ocean liner", "lipstick", "slip-on shoe", "lotion", "music speaker", "loupe magnifying glass", "sawmill", "magnetic compass", "messenger bag", "mailbox", "tights", "one-piece bathing suit", "manhole cover", "maraca", "marimba", "mask", "matchstick", "maypole", "maze", "measuring cup", "medicine cabinet", "megalith", "microphone", "microwave oven", "military uniform", "milk can", "minibus", "miniskirt", "minivan", "missile", "mitten", "mixing bowl", "mobile home", "ford model t", "modem", "monastery", "monitor", "moped", "mortar and pestle", "graduation cap", "mosque", "mosquito net", "vespa", "mountain bike", "tent", "computer mouse", "mousetrap", "moving van", "muzzle", "metal nail", "neck brace", "necklace", "baby pacifier", "notebook computer", "obelisk", "oboe", "ocarina", "odometer", "oil filter", "pipe organ", "oscilloscope", "overskirt", "bullock cart", "oxygen mask", "product packet / packaging", "paddle", "paddle wheel", "padlock", "paintbrush", "pajamas", "palace", "pan flute", "paper towel", "parachute", "parallel bars", "park bench", "parking meter", "railroad car", "patio", "payphone", "pedestal", "pencil case", "pencil sharpener", "perfume", "Petri dish", "photocopier", "plectrum", "Pickelhaube", "picket fence", "pickup truck", "pier", "piggy bank", "pill bottle", "pillow", "ping-pong ball", "pinwheel", "pirate ship", "drink pitcher", "block plane", "planetarium", "plastic bag", "plate rack", "farm plow", "plunger", "Polaroid camera", "pole", "police van", "poncho", "pool table", "soda bottle", "plant pot", "potter's wheel", "power drill", "prayer rug", "printer", "prison", "missile", "projector", "hockey puck", "punching bag", "purse", "quill", "quilt", "race car", "racket", "radiator", "radio", "radio telescope", "rain barrel", "recreational vehicle", "fishing casting reel", "reflex camera", "refrigerator", "remote control", "restaurant", "revolver", "rifle", "rocking chair", "rotisserie", "eraser", "rugby ball", "ruler measuring stick", "sneaker", "safe", "safety pin", "salt shaker", "sandal", "sarong", "saxophone", "scabbard", "weighing scale", "school bus", "schooner", "scoreboard", "CRT monitor", "screw", "screwdriver", "seat belt", "sewing machine", "shield", "shoe store", "shoji screen / room divider", "shopping basket", "shopping cart", "shovel", "shower cap", "shower curtain", "ski", "balaclava ski mask", "sleeping bag", "slide rule", "sliding door", "slot machine", "snorkel", "snowmobile", "snowplow", "soap dispenser", "soccer ball", "sock", "solar thermal collector", "sombrero", "soup bowl", "keyboard space bar", "space heater", "space shuttle", "spatula", "motorboat", "spider web", "spindle", "sports car", "spotlight", "stage", "steam locomotive", "through arch bridge", "steel drum", "stethoscope", "scarf", "stone wall", "stopwatch", "stove", "strainer", "tram", "stretcher", "couch", "stupa", "submarine", "suit", "sundial", "sunglasses", "sunglasses", "sunscreen", "suspension bridge", "mop", "sweatshirt", "swim trunks / shorts", "swing", "electrical switch", "syringe", "table lamp", "tank", "tape player", "teapot", "teddy bear", "television", "tennis ball", "thatched roof", "front curtain", "thimble", "threshing machine", "throne", "tile roof", "toaster", "tobacco shop", "toilet seat", "torch", "totem pole", "tow truck", "toy store", "tractor", "semi-trailer truck", "tray", "trench coat", "tricycle", "trimaran", "tripod", "triumphal arch", "trolleybus", "trombone", "hot tub", "turnstile", "typewriter keyboard", "umbrella", "unicycle", "upright piano", "vacuum cleaner", "vase", "vaulted or arched ceiling", "velvet fabric", "vending machine", "vestment", "viaduct", "violin", "volleyball", "waffle iron", "wall clock", "wallet", "wardrobe", "military aircraft", "sink", "washing machine", "water bottle", "water jug", "water tower", "whiskey jug", "whistle", "hair wig", "window screen", "window shade", "Windsor tie", "wine bottle", "airplane wing", "wok", "wooden spoon", "wool", "split-rail fence", "shipwreck", "sailboat", "yurt", "website", "comic book", "crossword", "traffic or street sign", "traffic light", "dust jacket", "menu", "plate", "guacamole", "consomme", "hot pot", "trifle", "ice cream", "popsicle", "baguette", "bagel", "pretzel", "cheeseburger", "hot dog", "mashed potatoes", "cabbage", "broccoli", "cauliflower", "zucchini", "spaghetti squash", "acorn squash", "butternut squash", "cucumber", "artichoke", "bell pepper", "cardoon", "mushroom", "Granny Smith apple", "strawberry", "orange", "lemon", "fig", "pineapple", "banana", "jackfruit", "cherimoya (custard apple)", "pomegranate", "hay", "carbonara", "chocolate syrup", "dough", "meatloaf", "pizza", "pot pie", "burrito", "red wine", "espresso", "tea cup", "eggnog", "mountain", "bubble", "cliff", "coral reef", "geyser", "lakeshore", "promontory", "sandbar", "beach", "valley", "volcano", "baseball player", "bridegroom", "scuba diver", "rapeseed", "daisy", "yellow lady's slipper", "corn", "acorn", "rose hip", "horse chestnut seed", "coral fungus", "agaric", "gyromitra", "stinkhorn mushroom", "earth star fungus", "hen of the woods mushroom", "bolete", "corn cob", "toilet paper"]

misc.py

@@ -0,0 +1,114 @@
+from itertools import repeat
+import collections.abc
+
+import torch
+from torch import nn as nn
+from torchvision.ops.misc import FrozenBatchNorm2d
+
+
+def freeze_batch_norm_2d(module, module_match={}, name=''):
+    """
+    Converts all `BatchNorm2d` and `SyncBatchNorm` layers of provided module into `FrozenBatchNorm2d`. If `module` is
+    itself an instance of either `BatchNorm2d` or `SyncBatchNorm`, it is converted into `FrozenBatchNorm2d` and
+    returned. Otherwise, the module is walked recursively and submodules are converted in place.
+
+    Args:
+        module (torch.nn.Module): Any PyTorch module.
+        module_match (dict): Dictionary of full module names to freeze (all if empty)
+        name (str): Full module name (prefix)
+
+    Returns:
+        torch.nn.Module: Resulting module
+
+    Inspired by https://github.com/pytorch/pytorch/blob/a5895f85be0f10212791145bfedc0261d364f103/torch/nn/modules/batchnorm.py#L762
+    """
+    res = module
+    is_match = True
+    if module_match:
+        is_match = name in module_match
+    if is_match and isinstance(module, (nn.modules.batchnorm.BatchNorm2d, nn.modules.batchnorm.SyncBatchNorm)):
+        res = FrozenBatchNorm2d(module.num_features)
+        res.num_features = module.num_features
+        res.affine = module.affine
+        if module.affine:
+            res.weight.data = module.weight.data.clone().detach()
+            res.bias.data = module.bias.data.clone().detach()
+        res.running_mean.data = module.running_mean.data
+        res.running_var.data = module.running_var.data
+        res.eps = module.eps
+    else:
+        for child_name, child in module.named_children():
+            full_child_name = '.'.join([name, child_name]) if name else child_name
+            new_child = freeze_batch_norm_2d(child, module_match, full_child_name)
+            if new_child is not child:
+                res.add_module(child_name, new_child)
+    return res
+
+
+# From PyTorch internals
+def _ntuple(n):
+    def parse(x):
+        if isinstance(x, collections.abc.Iterable):
+            return x
+        return tuple(repeat(x, n))
+    return parse
+
+
+to_1tuple = _ntuple(1)
+to_2tuple = _ntuple(2)
+to_3tuple = _ntuple(3)
+to_4tuple = _ntuple(4)
+to_ntuple = lambda n, x: _ntuple(n)(x)
+
+# Replaces all linear layers with linear_replacement
+# TODO: add int8 support for other linear layers including attn and convnets
+def replace_linear(model, linear_replacement, include_modules=['c_fc', 'c_proj'], copy_weights=True):
+    for name, module in model.named_children():
+        if len(list(module.children())) > 0:
+            replace_linear(module, linear_replacement, include_modules, copy_weights)
+
+        if isinstance(module, torch.nn.Linear) and name in include_modules:
+            old_module = model._modules[name]
+            model._modules[name] = linear_replacement(
+                module.in_features,
+                module.out_features,
+                module.bias is not None,
+            )
+            if copy_weights:
+                model._modules[name].weight.data.copy_(old_module.weight.data)
+                if model._modules[name].bias is not None:
+                    model._modules[name].bias.data.copy_(old_module.bias)
+
+    return model
+
+def convert_int8_model_to_inference_mode(model):
+    for m in model.modules():
+        if hasattr(m, 'prepare_for_eval'):
+            int8_original_dtype = m.weight.dtype
+            m.prepare_for_eval()
+            m.int8_original_dtype = int8_original_dtype
+            
+
+def accuracy(output, target, topk=(1,)):
+    """
+    Compute top-k accuracy
+
+    output: torch.Tensor
+        shape (N, C) where N is the number of examples, C the number of classes.
+        these are the logits.
+    
+    target: torch.Tensor
+        shape (N,) where N is the number of examples. Groundtruth class id of each example.
+    
+    topk: tuple
+        which topk to compute, e.g., topk=(1,5) will compute top-1 and top-5 accuracies
+    
+    Returns
+    -------
+    
+    list of top-k accuracies in the same order as `topk`
+    """
+    pred = output.topk(max(topk), 1, True, True)[1].t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+    n = len(target)
+    return [float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy()) / n for k in topk]

model.py

@@ -0,0 +1,431 @@
+""" CLIP Model
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+from dataclasses import dataclass
+import logging
+import math
+from typing import Optional, Tuple, Union, Text
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+
+from modified_resnet import ModifiedResNet
+from timm_model import TimmModel
+from transformer import LayerNorm, QuickGELU, VisionTransformer, TextTransformer, Attention
+from misc import to_2tuple
+
+
+
+@dataclass
+class CLIPVisionCfg:
+    layers: Union[Tuple[int, int, int, int], int] = 12
+    width: int = 768
+    head_width: int = 64
+    mlp_ratio: float = 4.0
+    patch_size: int = 16
+    image_size: Union[Tuple[int, int], int] = 224
+
+    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
+    input_patchnorm: bool = False  # whether to use dual patchnorm - would only apply the input layernorm on each patch, as post-layernorm already exist in original clip vit design
+    global_average_pool: bool = False  # whether to global average pool the last embedding layer, instead of using CLS token (https://arxiv.org/abs/2205.01580)
+    attentional_pool: bool = False  # whether to use attentional pooler in the last embedding layer
+    n_queries: int = 256  # n_queries for attentional pooler
+    attn_pooler_heads: int = 8  # n heads for attentional_pooling
+    output_tokens: bool = False
+
+    timm_model_name: str = None  # a valid model name overrides layers, width, patch_size
+    timm_model_pretrained: bool = False  # use (imagenet) pretrained weights for named model
+    timm_pool: str = 'avg'  # feature pooling for timm model ('abs_attn', 'rot_attn', 'avg', '')
+    timm_proj: str = 'linear'  # linear projection for timm model output ('linear', 'mlp', '')
+    timm_proj_bias: bool = False  # enable bias final projection
+    timm_drop: float = 0.  # head dropout
+    timm_drop_path: Optional[float] = None  # backbone stochastic depth
+
+
+
+
+def convert_weights_to_lp(model: nn.Module, dtype=torch.float16):
+    """Convert applicable model parameters to low-precision (bf16 or fp16)"""
+
+    def _convert_weights(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.to(dtype)
+            if l.bias is not None:
+                l.bias.data = l.bias.data.to(dtype)
+
+        if isinstance(l, (nn.MultiheadAttention, Attention)):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.to(dtype)
+
+        if isinstance(l, (CLIP, TextTransformer)):
+            # convert text nn.Parameter projections
+            attr = getattr(l, "text_projection", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+        if isinstance(l, VisionTransformer):
+            # convert vision nn.Parameter projections
+            attr = getattr(l, "proj", None)
+            if attr is not None:
+                attr.data = attr.data.to(dtype)
+
+    model.apply(_convert_weights)
+
+convert_weights_to_fp16 = convert_weights_to_lp  # backwards compat
+
+
+@dataclass
+class CLIPTextCfg:
+    context_length: int = 77
+    vocab_size: int = 49408
+    width: int = 512
+    heads: int = 8
+    layers: int = 12
+    ls_init_value: Optional[float] = None  # layer scale initial value
+    hf_model_name: str = None
+    hf_tokenizer_name: str = None
+    hf_model_pretrained: bool = True
+    proj: str = 'mlp'
+    pooler_type: str = 'mean_pooler'
+    embed_cls: bool = False
+    pad_id: int = 0
+    output_tokens: bool = False
+
+
+def get_cast_dtype(precision: str):
+    cast_dtype = None
+    if precision == 'bf16':
+        cast_dtype = torch.bfloat16
+    elif precision == 'fp16':
+        cast_dtype = torch.float16
+    return cast_dtype
+
+
+def get_input_dtype(precision: str):
+    input_dtype = None
+    if precision in ('bf16', 'pure_bf16'):
+        input_dtype = torch.bfloat16
+    elif precision in ('fp16', 'pure_fp16'):
+        input_dtype = torch.float16
+    return input_dtype
+
+
+def _build_vision_tower(
+        embed_dim: int,
+        vision_cfg: CLIPVisionCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(vision_cfg, dict):
+        vision_cfg = CLIPVisionCfg(**vision_cfg)
+
+    # OpenAI models are pretrained w/ QuickGELU but native nn.GELU is both faster and more
+    # memory efficient in recent PyTorch releases (>= 1.10).
+    # NOTE: timm models always use native GELU regardless of quick_gelu flag.
+    act_layer = QuickGELU if quick_gelu else nn.GELU
+
+    if vision_cfg.timm_model_name:
+        visual = TimmModel(
+            vision_cfg.timm_model_name,
+            pretrained=vision_cfg.timm_model_pretrained,
+            pool=vision_cfg.timm_pool,
+            proj=vision_cfg.timm_proj,
+            proj_bias=vision_cfg.timm_proj_bias,
+            drop=vision_cfg.timm_drop,
+            drop_path=vision_cfg.timm_drop_path,
+            patch_drop=vision_cfg.patch_dropout if vision_cfg.patch_dropout > 0 else None,
+            embed_dim=embed_dim,
+            image_size=vision_cfg.image_size,
+        )
+    elif isinstance(vision_cfg.layers, (tuple, list)):
+        vision_heads = vision_cfg.width * 32 // vision_cfg.head_width
+        visual = ModifiedResNet(
+            layers=vision_cfg.layers,
+            output_dim=embed_dim,
+            heads=vision_heads,
+            image_size=vision_cfg.image_size,
+            width=vision_cfg.width,
+        )
+    else:
+        vision_heads = vision_cfg.width // vision_cfg.head_width
+        norm_layer = LayerNorm
+        visual = VisionTransformer(
+            image_size=vision_cfg.image_size,
+            patch_size=vision_cfg.patch_size,
+            width=vision_cfg.width,
+            layers=vision_cfg.layers,
+            heads=vision_heads,
+            mlp_ratio=vision_cfg.mlp_ratio,
+            ls_init_value=vision_cfg.ls_init_value,
+            patch_dropout=vision_cfg.patch_dropout,
+            input_patchnorm=vision_cfg.input_patchnorm,
+            global_average_pool=vision_cfg.global_average_pool,
+            attentional_pool=vision_cfg.attentional_pool,
+            n_queries=vision_cfg.n_queries,
+            attn_pooler_heads=vision_cfg.attn_pooler_heads,
+            output_tokens=vision_cfg.output_tokens,
+            output_dim=embed_dim,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+    return visual
+
+
+def _build_text_tower(
+        embed_dim: int,
+        text_cfg: CLIPTextCfg,
+        quick_gelu: bool = False,
+        cast_dtype: Optional[torch.dtype] = None,
+):
+    if isinstance(text_cfg, dict):
+        text_cfg = CLIPTextCfg(**text_cfg)
+
+    if text_cfg.hf_model_name:
+        from hf_model import HFTextEncoder
+        text = HFTextEncoder(
+            text_cfg.hf_model_name,
+            output_dim=embed_dim,
+            proj=text_cfg.proj,
+            pooler_type=text_cfg.pooler_type,
+            pretrained=text_cfg.hf_model_pretrained,
+            output_tokens=text_cfg.output_tokens,
+        )
+    else:
+        act_layer = QuickGELU if quick_gelu else nn.GELU
+        norm_layer = LayerNorm
+
+        text = TextTransformer(
+            context_length=text_cfg.context_length,
+            vocab_size=text_cfg.vocab_size,
+            width=text_cfg.width,
+            heads=text_cfg.heads,
+            layers=text_cfg.layers,
+            ls_init_value=text_cfg.ls_init_value,
+            output_dim=embed_dim,
+            embed_cls=text_cfg.embed_cls,
+            output_tokens=text_cfg.output_tokens,
+            pad_id=text_cfg.pad_id,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+    return text
+
+
+class CLIP(nn.Module):
+    """
+        _VITL14_336 = dict(
+        openai=_pcfg(
+            url="https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+            hf_hub="timm/vit_large_patch14_clip_336.openai/",
+            quick_gelu=True,
+        ),
+    )
+    
+    """
+    output_dict: torch.jit.Final[bool]
+
+    def __init__(
+            self,
+            embed_dim: int,
+            vision_cfg: CLIPVisionCfg,
+            text_cfg: CLIPTextCfg,
+            quick_gelu: bool = False,
+            cast_dtype: Optional[torch.dtype] = None,
+            output_dict: bool = False,
+    ):
+        super().__init__()
+        self.output_dict = output_dict
+        self.visual = _build_vision_tower(embed_dim, vision_cfg, quick_gelu, cast_dtype)
+        print(f"Building vision tower with config: {vision_cfg}")
+        
+        print(f"Currently text tower is removed, using only image encoder for feature extraction")
+        do_use = False
+        if do_use:
+            text = _build_text_tower(embed_dim, text_cfg, quick_gelu, cast_dtype)
+            self.transformer = text.transformer
+            self.context_length = text.context_length
+            self.vocab_size = text.vocab_size
+            self.token_embedding = text.token_embedding
+            self.positional_embedding = text.positional_embedding
+            self.ln_final = text.ln_final
+            self.text_projection = text.text_projection
+            self.register_buffer('attn_mask', text.attn_mask, persistent=False)
+
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.num_register_tokens = 0
+        self.neuron_dict = None
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.visual.set_grad_checkpointing(enable)
+        self.transformer.grad_checkpointing = enable
+
+    def encode_image(self, image, normalize: bool = False, attn_method: Text = 'direct', num_register_tokens = None, neuron_dict=None):
+        if num_register_tokens is None and neuron_dict is None:
+            num_register_tokens = self.num_register_tokens
+            neuron_dict = self.neuron_dict
+            
+        
+        features = self.visual(image, attn_method=attn_method, num_register_tokens=num_register_tokens, neuron_dict=neuron_dict)
+        return F.normalize(features, dim=-1) if normalize else features
+
+    def encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        # x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        # x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+        return F.normalize(x, dim=-1) if normalize else x
+
+    def forward(
+            self,
+            image: Optional[torch.Tensor] = None,
+            text: Optional[torch.Tensor] = None,
+            num_register_tokens = None,
+            neuron_dict=None
+
+    ):
+        
+        if num_register_tokens is None and neuron_dict is None:
+            num_register_tokens = self.num_register_tokens
+            neuron_dict = self.neuron_dict
+
+
+        image_features = self.encode_image(image, num_register_tokens=num_register_tokens, neuron_dict=neuron_dict, normalize=True) if image is not None else None
+        text_features = self.encode_text(text, normalize=True) if text is not None else None
+        if self.output_dict:
+            return {
+                "image_features": image_features,
+                "text_features": text_features,
+                "logit_scale": self.logit_scale.exp()
+            }
+        return image_features, text_features, self.logit_scale.exp()
+
+
+# used to maintain checkpoint compatibility
+def convert_to_custom_text_state_dict(state_dict: dict):
+    if 'text_projection' in state_dict:
+        # old format state_dict, move text tower -> .text
+        new_state_dict = {}
+        for k, v in state_dict.items():
+            if any(k.startswith(p) for p in (
+                'text_projection',
+                'positional_embedding',
+                'token_embedding',
+                'transformer',
+                'ln_final',
+            )):
+                k = 'text.' + k
+            new_state_dict[k] = v
+        return new_state_dict
+    return state_dict
+
+
+def build_model_from_openai_state_dict(
+        state_dict: dict,
+        quick_gelu=True,
+        cast_dtype=torch.float16,
+):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len(
+            [k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_size = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_size = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    vision_cfg = CLIPVisionCfg(
+        layers=vision_layers,
+        width=vision_width,
+        patch_size=vision_patch_size,
+        image_size=image_size,
+    )
+    text_cfg = CLIPTextCfg(
+        context_length=context_length,
+        vocab_size=vocab_size,
+        width=transformer_width,
+        heads=transformer_heads,
+        layers=transformer_layers,
+    )
+    model = CLIP(
+        embed_dim,
+        vision_cfg=vision_cfg,
+        text_cfg=text_cfg,
+        quick_gelu=quick_gelu,  # OpenAI models were trained with QuickGELU
+        cast_dtype=cast_dtype,
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        state_dict.pop(key, None)
+
+    convert_weights_to_fp16(model)  # OpenAI state dicts are partially converted to float16
+    model.load_state_dict(state_dict)
+    return model.eval()
+
+
+def resize_pos_embed(state_dict, model, interpolation: str = 'bicubic', antialias: bool = True):
+    # Rescale the grid of position embeddings when loading from state_dict
+    old_pos_embed = state_dict.get('visual.positional_embedding', None)
+    if old_pos_embed is None or not hasattr(model.visual, 'grid_size'):
+        return
+    grid_size = to_2tuple(model.visual.grid_size)
+    extra_tokens = 1  # FIXME detect different token configs (ie no class token, or more)
+    new_seq_len = grid_size[0] * grid_size[1] + extra_tokens
+    if new_seq_len == old_pos_embed.shape[0]:
+        return
+
+    if extra_tokens:
+        pos_emb_tok, pos_emb_img = old_pos_embed[:extra_tokens], old_pos_embed[extra_tokens:]
+    else:
+        pos_emb_tok, pos_emb_img = None, old_pos_embed
+    old_grid_size = to_2tuple(int(math.sqrt(len(pos_emb_img))))
+
+    logging.info('Resizing position embedding grid-size from %s to %s', old_grid_size, grid_size)
+    pos_emb_img = pos_emb_img.reshape(1, old_grid_size[0], old_grid_size[1], -1).permute(0, 3, 1, 2)
+    pos_emb_img = F.interpolate(
+        pos_emb_img,
+        size=grid_size,
+        mode=interpolation,
+        antialias=antialias,
+        align_corners=False,
+    )
+    pos_emb_img = pos_emb_img.permute(0, 2, 3, 1).reshape(1, grid_size[0] * grid_size[1], -1)[0]
+    if pos_emb_tok is not None:
+        new_pos_embed = torch.cat([pos_emb_tok, pos_emb_img], dim=0)
+    else:
+        new_pos_embed = pos_emb_img
+    state_dict['visual.positional_embedding'] = new_pos_embed

model_sanity_check.ipynb

@@ -0,0 +1,285 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "ba945813",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "e7cec94e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "import os, json, math, torch, tqdm\n",
+    "from pathlib import Path\n",
+    "from torchvision import transforms\n",
+    "from torchvision.datasets import ImageFolder\n",
+    "from torch.utils.data import DataLoader\n",
+    "from transformers import CLIPProcessor, CLIPModel\n",
+    "import os\n",
+    "import itertools\n",
+    "\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "\n",
+    "import transformers\n",
+    "from transformers import AutoModel, AutoProcessor, CLIPForImageClassification, AutoConfig, AutoTokenizer\n",
+    "from torchvision import transforms\n",
+    "from torchvision.datasets import ImageNet\n",
+    "from torch.utils.data import Subset\n",
+    "from tqdm import tqdm\n",
+    "from PIL import Image\n",
+    "import matplotlib.ticker as mticker\n",
+    "import matplotlib.pyplot as plt\n",
+    "from mpl_toolkits.mplot3d import Axes3D   # noqa: F401 – 3D 기능 활성화\n",
+    "import inspect\n",
+    "import torch.nn.functional as F\n",
+    "import torchvision.transforms.functional as VF\n",
+    "import tqdm\n",
+    "\n",
+    "\n",
+    "from functools import partial\n",
+    "\n",
+    "from torchvision import transforms\n",
+    "from torchvision.transforms import InterpolationMode\n",
+    "\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import yaml\n",
+    "from pathlib import Path\n",
+    "\n",
+    "import sys\n",
+    "import os\n",
+    "\n",
+    "from imagenet_classes import *\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4c7a750",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Pretrained path from config: /workspace/code/clipL336_TTR\n",
+      "✓ Added '/workspace/code/clipL336_TTR' to Python path.\n",
+      "✓ Successfully imported 'model' from '/workspace/code/clipL336_TTR'\n",
+      "Building vision tower with config: CLIPVisionCfg(layers=24, width=1024, head_width=64, mlp_ratio=4.0, patch_size=14, image_size=336, ls_init_value=None, patch_dropout=0.0, input_patchnorm=False, global_average_pool=False, attentional_pool=False, n_queries=256, attn_pooler_heads=8, output_tokens=False, timm_model_name=None, timm_model_pretrained=False, timm_pool='avg', timm_proj='linear', timm_proj_bias=False, timm_drop=0.0, timm_drop_path=None)\n",
+      "✓ Added '/workspace/data/cache/huggingface/modules/transformers_modules/clipL336_TTR' to Python path.\n",
+      "✓ Successfully imported 'tokenizer' from '/workspace/data/cache/huggingface/modules/transformers_modules/clipL336_TTR'\n",
+      "Custom CLIP model loaded successfully!\n"
+     ]
+    }
+   ],
+   "source": [
+    "# 문제의 원인이 text encoder가 고장이 나 있었다..\n",
+    "device = \"cuda:7\"\n",
+    "model_path = \"/workspace/code/clipL336_TTR\"\n",
+    "\n",
+    "exp_cfg = AutoConfig.from_pretrained(\"/workspace/code/clipL336_TTR\", trust_remote_code=True)\n",
+    "exp_cfg.pretrained_path = model_path \n",
+    "# model = AutoModel.from_pretrained(model_path, trust_remote_code=True, local_files_only=True)\n",
+    "model = AutoModel.from_pretrained(pretrained_model_name_or_path=model_path, config=exp_cfg, trust_remote_code=True, local_files_only=True)\n",
+    "# 여기 load 되었는 지 확인할 필요 있음\n",
+    "model = model.to(device)\n",
+    "preprocessor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True, local_files_only=True)\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, local_files_only=True)\n",
+    "# tokenizer랑 preprocessor 가져오기\n",
+    "\n",
+    "clip_transform = lambda image: preprocessor.image_processor(image, return_tensors=\"pt\")['pixel_values'].squeeze(0)  # 와 이렇게 활용할 방법은 생각도 못했네\n",
+    "model_clip = AutoModel.from_pretrained(\"openai/clip-vit-large-patch14-336\").to(device).half()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "ed3cbfdc",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 1000/1000 [00:23<00:00, 41.71it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Built text features: torch.Size([768, 1000])\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "# langauge head\n",
+    "### zeroshot head construction (text encoding) ###\n",
+    "with torch.no_grad():\n",
+    "    zeroshot_weight = []\n",
+    "    for classname in tqdm(IMAGENET_CLASSNAMES):\n",
+    "        texts = [template(classname) for template in OPENAI_IMAGENET_TEMPLATES]\n",
+    "        text_inputs = preprocessor(text=texts, return_tensors=\"pt\", padding=\"max_length\").to(device)\n",
+    "        # text_inputs = model.tokenize(texts).to(device)\n",
+    "        # text_features = model.encode_text(text_inputs.input_ids)\n",
+    "        text_features = model_clip.get_text_features(**text_inputs)\n",
+    "        text_feature = F.normalize(text_features, dim=-1).mean(dim=0)\n",
+    "        # text_feature = text_features.mean(dim=0)\n",
+    "        text_feature = text_feature / text_feature.norm()\n",
+    "        zeroshot_weight.append(text_feature)\n",
+    "    \n",
+    "    text_features = torch.stack(zeroshot_weight, dim=1).to(device)\n",
+    "print(\"Built text features:\", text_features.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "e1bd37d1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch.save(text_features, \"./zeroshot_classifier.pt\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "dbfeaedf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "imagenet_dataset = ImageNet(root='/workspace/data/imagenet', split='val', transform=clip_transform)\n",
+    "eval_loader = torch.utils.data.DataLoader(imagenet_dataset, batch_size=128, num_workers=16, pin_memory=False, shuffle=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "b0000195",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "\n",
+    "half = \"torch.bfloat16\"\n",
+    "def evaluate(model, loader, text_feats, max_samples: int | None = None):\n",
+    "    model.eval()\n",
+    "    top1 = top5 = n = 0\n",
+    "    pbar = tqdm(loader, desc=\"Evaluating\", unit=\"batch\")\n",
+    "    with torch.no_grad():\n",
+    "        for images, labels in pbar:\n",
+    "            if max_samples and n >= max_samples:\n",
+    "                break\n",
+    "            images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True)\n",
+    "            with torch.autocast(device_type=\"cuda\"):\n",
+    "                # 여기 test-time 가공 함수 구현 필요\n",
+    "                feats = model.encode_image(images)\n",
+    "\n",
+    "                feats = feats / feats.norm(dim=-1, keepdim=True)\n",
+    "                logits = model.model.logit_scale.exp() * feats @ text_feats \n",
+    "            _, pred = logits.topk(5, dim=-1)\n",
+    "            top1 += (pred[:, :1] == labels.unsqueeze(1)).sum().item()\n",
+    "            top5 += (pred == labels.unsqueeze(1)).sum().item()\n",
+    "            n += images.size(0)\n",
+    "            pbar.set_postfix(samples=n, top1=top1 / n * 100, top5=top5 / n * 100)\n",
+    "        return top1 / n * 100, top5 / n * 100\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "8795b394",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Evaluating:   0%|          | 0/391 [00:00<?, ?batch/s]"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Evaluating: 100%|██████████| 391/391 [10:38<00:00,  1.63s/batch, samples=5e+4, top1=74.9, top5=94.4] "
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Baseline (Top‑1 / Top‑5) on 50,000 imgs: 74.87% / 94.37%\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "\n",
+    "### baseline evaluator ###\n",
+    "### 이거는 지금 당장은 못 써먹는다... 미친 너무 느리다 어디서 문제지 ###\n",
+    "# 씨발 이번에 뭐지\n",
+    "# architecture define이 어딘가에서 손상 된 것으로 보인다\n",
+    "# 성능 reproduce...\n",
+    "\n",
+    "BASELINE_SAMPLES = 50000  # set to None for full 50 k\n",
+    "acc1, acc5 = evaluate(model, eval_loader, text_features, max_samples=BASELINE_SAMPLES)\n",
+    "print(f\"Baseline (Top‑1 / Top‑5) on {BASELINE_SAMPLES or len(imagenet_dataset):,} imgs: {acc1:.2f}% / {acc5:.2f}%\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4aa82bb4",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

modeling_custom_clip.py

@@ -0,0 +1,357 @@
+"""
+Custom CLIP Model with Register Tokens - Import Safe Version with Complete File Download
+"""
+import transformers
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.utils import logging
+from typing import Optional, Union, Tuple
+import json
+from pathlib import Path
+import warnings
+import os
+import sys
+import importlib.util
+
+# Suppress all warnings during import
+warnings.filterwarnings("ignore")
+os.environ["TRANSFORMERS_VERBOSITY"] = "error"
+
+logger = logging.get_logger(__name__)
+
+
+def safe_import_from_repo(module_name: str, repo_path: str):
+    """
+    지정된 로컬 경로(repo_path)에서 파이썬 모듈을 안전하게 임포트합니다.
+
+    Args:
+        module_name (str): 임포트할 모듈의 이름 (예: 'modeling_clip').
+        repo_path (str): 모듈이 포함된 디렉토리의 경로.
+
+    Returns:
+        The imported module object.
+
+    Raises:
+        ValueError: repo_path가 None이거나 유효한 디렉토리가 아닐 경우.
+        ImportError: 지정된 경로에서 모듈을 찾을 수 없을 경우.
+    """
+    # 1. repo_path가 유효한지 검사합니다.
+    if repo_path is None:
+        raise ValueError("The 'repo_path' argument cannot be None.")
+
+    # pathlib.Path 객체로 변환하여 경로를 쉽게 다룰 수 있도록 합니다.
+    repo_path_obj = Path(repo_path)
+
+    if not repo_path_obj.is_dir():
+        raise ValueError(
+            f"The specified repo_path does not exist or is not a directory: '{repo_path}'")
+
+    # 2. 파이썬이 모듈을 찾을 수 있도록 해당 경로를 sys.path에 추가합니다.
+    #    resolve()를 통해 절대 경로를 사용하고, 문자열로 변환합니다.
+    absolute_repo_path = str(repo_path_obj.resolve())
+
+    if absolute_repo_path not in sys.path:
+        # sys.path의 맨 앞에 추가하여 다른 경로보다 우선적으로 탐색되도록 합니다.
+        sys.path.insert(0, absolute_repo_path)
+        print(f"✓ Added '{absolute_repo_path}' to Python path.")
+
+    # 3. `importlib`을 사용하여 모듈을 동적으로 임포트합니다.
+    try:
+        module = importlib.import_module(module_name)
+        print(
+            f"✓ Successfully imported '{module_name}' from '{absolute_repo_path}'")
+        return module
+    except ImportError:
+        # sys.path에 경로를 추가했음에도 임포트에 실패한 경우,
+        # 해당 경로에 모듈 파일(.py)이 없다는 의미입니다.
+        raise ImportError(
+            f"Module '{module_name}' not found inside the specified path: '{absolute_repo_path}'")
+
+
+class CustomCLIPConfig(PretrainedConfig):
+    model_type = "custom_clip_with_registers"
+
+    def __init__(
+        self,
+        vision_config=None,
+        text_config=None,
+        num_register_tokens=0,
+        neuron_dict=None,
+        projection_dim=512,
+        logit_scale_init_value=2.6592,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        self.vision_config = vision_config or {}
+        self.text_config = text_config or {}
+        self.num_register_tokens = num_register_tokens
+        self.neuron_dict = neuron_dict
+        self.projection_dim = projection_dim
+        self.logit_scale_init_value = logit_scale_init_value
+
+
+class CustomCLIPModel(PreTrainedModel):
+    config_class = CustomCLIPConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        # Safe import of custom modules
+        try:
+            # to strictly load from the local library
+            pretrained_path: str | None = getattr(
+                config, "pretrained_path", None)
+            if pretrained_path is None:
+                raise ValueError(
+                    "The config must have a 'pretrained_path' attribute pointing to the local repository path.")
+            else:
+                print(f"Pretrained path from config: {pretrained_path}")
+
+            model_module = safe_import_from_repo('model', pretrained_path)
+            self.CLIP = model_module.CLIP
+            self.CLIPVisionCfg = model_module.CLIPVisionCfg
+            self.CLIPTextCfg = model_module.CLIPTextCfg
+        except ImportError as e:
+            raise ImportError(
+                f"Could not import model components: {e}. Make sure all model files are in the repository.")
+
+        # Create vision and text configs
+        vision_cfg = self.CLIPVisionCfg(
+            layers=config.vision_config.get("num_hidden_layers", 12),
+            width=config.vision_config.get("hidden_size", 768),
+            patch_size=config.vision_config.get("patch_size", 16),
+            image_size=config.vision_config.get("image_size", 224),
+        )
+
+        text_cfg = self.CLIPTextCfg(
+            context_length=config.text_config.get(
+                "max_position_embeddings", 77),
+            vocab_size=config.text_config.get("vocab_size", 49408),
+            width=config.text_config.get("hidden_size", 512),
+            layers=config.text_config.get("num_hidden_layers", 12),
+        )
+
+        # Initialize your custom CLIP model
+        self.model = self.CLIP(
+            embed_dim=config.projection_dim,
+            vision_cfg=vision_cfg,
+            text_cfg=text_cfg,
+        )
+
+        # These will be set when loading the state dict
+        # 여기 statedict에서 load하면 않된다. configuration에 떡하니 있으면서 무슨 짓거리냐
+        self.neuron_dict = config.neuron_dict
+        if self.neuron_dict is None:
+            raise ValueError("neuron_dict must be provided in the config.")
+        self.num_register_tokens = config.num_register_tokens
+
+        # These will be loaded separately
+        self._tokenizer = None
+        self._preprocessor = None
+        self._zeroshot_classifier = None
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        """Override to handle custom parameters and load weights properly"""
+
+        # Extract custom parameters first
+        if 'neuron_dict' in state_dict:
+            self.neuron_dict = state_dict.pop('neuron_dict')
+
+        if 'num_register_tokens' in state_dict:
+            self.num_register_tokens = state_dict.pop('num_register_tokens')
+
+        # Set these values in the model
+        if hasattr(self.model, 'visual'):
+            self.model.visual.num_register_tokens = self.num_register_tokens
+            self.model.visual.neuron_dict = self.neuron_dict
+            self.model.num_register_tokens = self.num_register_tokens
+            self.model.neuron_dict = self.neuron_dict
+
+        # Load the weights properly - suppress ALL warnings and errors
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+
+            # Temporarily set logging to critical only
+            original_level = logging.get_verbosity()
+            logging.set_verbosity_error()
+
+            try:
+                # Load weights directly into self.model
+                missing, unexpected = self.model.load_state_dict(
+                    state_dict, strict=False)
+
+                # Don't report any missing/unexpected keys to avoid warnings
+
+            except Exception as e:
+                # If direct loading fails, try the parent method silently
+                super()._load_from_state_dict(state_dict, prefix, local_metadata, False, [], [], [])
+            finally:
+                # Restore logging level
+                logging.set_verbosity(original_level)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        """Override to load cleanly and suppress warnings"""
+
+        # Suppress warnings during loading
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+
+            # Temporarily suppress transformers logging
+            original_level = logging.get_verbosity()
+            logging.set_verbosity_error()
+
+            try:
+                # Load the model
+                model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+            finally:
+                # Restore logging
+                logging.set_verbosity(original_level)
+
+        # Load additional components
+        model._load_additional_components(pretrained_model_name_or_path)
+
+        # Print clean success message
+        print("Custom CLIP model loaded successfully!")
+
+        return model
+
+    def _load_additional_components(self, pretrained_model_name_or_path):
+        """Load tokenizer, preprocessor, and zero-shot classifier silently"""
+
+        try:
+            from huggingface_hub import hf_hub_download
+
+            # Load tokenizer
+            try:
+                # Safe import of tokenizer
+                tokenizer_module = safe_import_from_repo(
+                    'tokenizer', Path(__file__).parent)
+                self._tokenizer = tokenizer_module.SimpleTokenizer()
+            except ImportError:
+                # If tokenizer import fails, create a dummy tokenizer message
+                pass
+
+            # Load preprocessor
+            try:
+                preprocess_config_file = hf_hub_download(
+                    repo_id=pretrained_model_name_or_path,
+                    filename="preprocessor_config.json"
+                )
+
+                with open(preprocess_config_file, 'r') as f:
+                    preprocess_config = json.load(f)
+
+                self._create_preprocessor(preprocess_config)
+            except:
+                pass
+
+            # Load zero-shot classifier
+            try:
+                classifier_file = hf_hub_download(
+                    repo_id=pretrained_model_name_or_path,
+                    filename="zeroshot_classifier.pt"
+                )
+
+                # Suppress the torch.load warning
+                with warnings.catch_warnings():
+                    warnings.simplefilter("ignore")
+                    self._zeroshot_classifier = torch.load(
+                        classifier_file, map_location='cpu', weights_only=False)
+            except:
+                pass
+
+        except:
+            pass
+
+    def _create_preprocessor(self, config):
+        """Create image preprocessor from config"""
+        try:
+            from torchvision import transforms
+
+            self._preprocessor = transforms.Compose([
+                transforms.Resize(
+                    config["image_size"], interpolation=transforms.InterpolationMode.BICUBIC),
+                transforms.CenterCrop(config["image_size"]),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=config["image_mean"], std=config["image_std"]),
+            ])
+        except:
+            pass
+
+    @property
+    def tokenizer(self):
+        """Access the tokenizer"""
+        return self._tokenizer
+
+    @property
+    def preprocessor(self):
+        """Access the image preprocessor"""
+        return self._preprocessor
+
+    @property
+    def zeroshot_classifier(self):
+        """Access the zero-shot classifier"""
+        return self._zeroshot_classifier
+
+    def tokenize(self, texts, context_length=77):
+        """Tokenize text using the loaded tokenizer"""
+        if self._tokenizer is None:
+            raise ValueError(
+                "Tokenizer not available. Make sure tokenizer.py is in the repository.")
+
+        # Safe import of tokenize function
+        try:
+            tokenizer_module = safe_import_from_repo(
+                'tokenizer', Path(__file__).parent)
+            return tokenizer_module.tokenize(texts, context_length)
+        except ImportError:
+            raise ValueError("Could not import tokenize function.")
+
+    def preprocess_image(self, image):
+        """Preprocess image using the loaded preprocessor"""
+        if self._preprocessor is None:
+            raise ValueError(
+                "Preprocessor not loaded. Make sure preprocessor_config.json is in the repository.")
+
+        return self._preprocessor(image)
+
+    def forward(self, input_ids=None, pixel_values=None, num_register_tokens=None, neuron_dict=None, **kwargs):
+        """Forward pass supporting your custom functionality"""
+
+        if num_register_tokens is None:
+            num_register_tokens = self.num_register_tokens
+        if neuron_dict is None:
+            neuron_dict = self.neuron_dict
+
+        return self.model(
+            image=pixel_values,
+            text=input_ids,
+            num_register_tokens=num_register_tokens,
+            neuron_dict=neuron_dict
+        )
+
+    def encode_image(self, pixel_values, num_register_tokens=None, neuron_dict=None, **kwargs):
+        """Encode images with register token support"""
+        if num_register_tokens is None:
+            num_register_tokens = self.num_register_tokens
+        if neuron_dict is None:
+            neuron_dict = self.neuron_dict
+
+        return self.model.encode_image(
+            pixel_values,
+            num_register_tokens=num_register_tokens,
+            neuron_dict=neuron_dict,
+            **kwargs
+        )
+
+    def encode_text(self, input_ids, **kwargs):
+        """Encode text"""
+        return self.model.encode_text(input_ids, **kwargs)
+
+
+# Auto-suppress warnings at module level
+transformers.logging.set_verbosity_error()

modified_resnet.py

@@ -0,0 +1,181 @@
+from collections import OrderedDict
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from misc import freeze_batch_norm_2d
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.act1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.act2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.act3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.act1(self.bn1(self.conv1(x)))
+        out = self.act2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.act3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x):
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x, key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0.,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+
+        return x[0]
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, image_size=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.image_size = image_size
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.act1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.act2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.act3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(image_size // 32, embed_dim, heads, output_dim)
+
+        self.init_parameters()
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def init_parameters(self):
+        if self.attnpool is not None:
+            std = self.attnpool.c_proj.in_features ** -0.5
+            nn.init.normal_(self.attnpool.q_proj.weight, std=std)
+            nn.init.normal_(self.attnpool.k_proj.weight, std=std)
+            nn.init.normal_(self.attnpool.v_proj.weight, std=std)
+            nn.init.normal_(self.attnpool.c_proj.weight, std=std)
+
+        for resnet_block in [self.layer1, self.layer2, self.layer3, self.layer4]:
+            for name, param in resnet_block.named_parameters():
+                if name.endswith("bn3.weight"):
+                    nn.init.zeros_(param)
+
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        assert unlocked_groups == 0, 'partial locking not currently supported for this model'
+        for param in self.parameters():
+            param.requires_grad = False
+        if freeze_bn_stats:
+            freeze_batch_norm_2d(self)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        # FIXME support for non-transformer
+        pass
+
+    def stem(self, x):
+        x = self.act1(self.bn1(self.conv1(x)))
+        x = self.act2(self.bn2(self.conv2(x)))
+        x = self.act3(self.bn3(self.conv3(x)))
+        x = self.avgpool(x)
+        return x
+
+    def forward(self, x):
+        x = self.stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.attnpool(x)
+
+        return x

neuron_indices.json

@@ -0,0 +1 @@
+[[12, 42, 39.99140930175781], [12, 983, 34.50058364868164], [12, 3868, 23.993741989135742], [12, 2687, 23.192779541015625], [11, 3784, 14.847213745117188], [11, 987, 14.675474166870117], [11, 3661, 14.301347732543945], [12, 3008, 12.25265884399414], [11, 1967, 11.993508338928223], [12, 3002, 10.681584358215332], [11, 9, 9.61478042602539], [21, 1801, 8.448626518249512], [11, 2555, 6.903197288513184], [11, 1100, 6.859874725341797], [12, 1571, 4.70828104019165], [22, 901, 3.453416109085083], [21, 1550, 3.4134912490844727], [12, 1816, 3.37734055519104], [12, 183, 3.1418349742889404], [8, 745, 3.1221530437469482], [9, 4078, 3.0656824111938477], [9, 815, 3.0607407093048096], [10, 357, 2.7818374633789062], [9, 3618, 2.690423011779785], [10, 1654, 2.6796107292175293], [22, 2184, 2.6561291217803955], [10, 3940, 2.652881383895874], [7, 3228, 2.46209979057312], [10, 2319, 2.308473825454712], [9, 2693, 2.1979129314422607], [21, 1779, 2.1429498195648193], [20, 3077, 2.1137425899505615], [20, 2634, 2.04282808303833], [9, 3973, 2.031193733215332], [21, 3137, 2.026745080947876], [8, 3249, 1.9856672286987305], [8, 2585, 1.9620095491409302], [9, 1983, 1.9459211826324463], [9, 1744, 1.9378128051757812], [9, 1157, 1.749971866607666], [21, 2412, 1.7358660697937012], [10, 2560, 1.6931447982788086], [7, 2550, 1.6547895669937134], [21, 1381, 1.5941085815429688], [22, 1317, 1.560852289199829], [8, 1537, 1.5494486093521118], [8, 200, 1.4573794603347778], [19, 1881, 1.4518368244171143], [8, 1603, 1.416003704071045], [8, 1851, 1.3301061391830444], [8, 3523, 1.321004867553711], [12, 2780, 1.2789242267608643], [13, 1109, 1.2571412324905396], [10, 2559, 1.2549676895141602], [9, 1309, 1.238487958908081], [21, 2193, 1.2044764757156372], [17, 1868, 1.1777989864349365], [21, 1796, 1.1429805755615234], [10, 4009, 1.0898690223693848], [9, 1335, 1.0648274421691895], [22, 2889, 1.0604228973388672], [11, 888, 1.0271515846252441], [15, 415, 1.0207806825637817], [21, 68, 1.0149273872375488], [9, 3049, 0.9941853880882263], [9, 2607, 0.9631124138832092], [9, 2621, 0.954177737236023], [18, 1283, 0.9397207498550415], [9, 2396, 0.9153910875320435], [22, 797, 0.8976885676383972], [12, 2370, 0.8916781544685364], [22, 3026, 0.8911128044128418], [10, 3029, 0.864679217338562], [19, 2881, 0.8607441782951355], [9, 1610, 0.8600460886955261], [22, 3143, 0.8545671105384827], [19, 1149, 0.8446468114852905], [22, 806, 0.8359670042991638], [20, 676, 0.8346958756446838], [18, 3018, 0.8332728147506714], [22, 2714, 0.8295024037361145], [9, 2867, 0.813927948474884], [22, 3888, 0.8113453388214111], [8, 3697, 0.8057175278663635], [22, 1832, 0.7937105894088745], [22, 985, 0.7906701564788818], [22, 3361, 0.783061683177948], [9, 2394, 0.7818043231964111], [22, 3049, 0.7765958309173584], [8, 3137, 0.772114098072052], [10, 951, 0.7676676511764526], [11, 3568, 0.7665989398956299], [8, 2563, 0.7626394629478455], [23, 1137, 0.7513239979743958], [17, 604, 0.7489021420478821], [9, 1924, 0.7423470616340637], [19, 2106, 0.7369621992111206], [9, 2012, 0.7241123914718628], [10, 1903, 0.7238287329673767], [12, 3574, 0.7192695736885071]]

openai_models.py

@@ -0,0 +1,91 @@
+""" OpenAI pretrained model functions
+
+Adapted from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+
+import os
+import warnings
+from typing import List, Optional, Union
+
+import torch
+
+from constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+from model import build_model_from_openai_state_dict, get_cast_dtype
+from pretrained import *
+
+__all__ = ["list_openai_models", "load_openai_model"]
+
+
+def list_openai_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list_pretrained_models_by_tag('openai')
+
+
+def load_openai_model(
+        name: str,
+        precision: Optional[str] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        cache_dir: Optional[str] = None,
+        quick_gelu: Optional[bool] = True
+):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+    precision: str
+        Model precision, if None defaults to 'fp32' if device == 'cpu' else 'fp16'.
+    device : Union[str, torch.device]
+        The device to put the loaded model
+    cache_dir : Optional[str]
+        The directory to cache the downloaded model weights
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    if precision is None:
+        precision = 'fp32' if device == 'cpu' else 'fp16'
+
+    if get_pretrained_url(name, 'openai'):
+        model_path = download_pretrained_from_url(get_pretrained_url(name, 'openai'), cache_dir=cache_dir)
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {list_openai_models()}")
+
+    try:
+        # loading JIT archive
+        model = torch.jit.load(model_path, map_location="cpu").eval()
+        state_dict = None
+    except RuntimeError:
+        # loading saved state dict
+        state_dict = torch.load(model_path, map_location="cpu")
+
+    # Build a non-jit model from the OpenAI jitted model state dict
+    cast_dtype = get_cast_dtype(precision)
+    try:
+        model = build_model_from_openai_state_dict(state_dict or model.state_dict(), quick_gelu=quick_gelu, cast_dtype=cast_dtype)
+    except KeyError:
+        sd = {k[7:]: v for k, v in state_dict["state_dict"].items()}
+        model = build_model_from_openai_state_dict(sd, cast_dtype=cast_dtype, quick_gelu=quick_gelu)
+
+    # model from OpenAI state dict is in manually cast fp16 mode, must be converted for AMP/fp32/bf16 use
+    model = model.to(device)
+    # FIXME support pure fp16/bf16 precision modes
+    if precision != 'fp16':
+        model.float()
+        if precision == 'bf16':
+            # for bf16, convert back to low-precision
+            convert_weights_to_lp(model, dtype=torch.bfloat16)
+
+    # add mean / std attributes for consistency with OpenCLIP models
+    model.visual.image_mean = OPENAI_DATASET_MEAN
+    model.visual.image_std = OPENAI_DATASET_STD
+    return model

openai_templates.py

@@ -0,0 +1,84 @@
+
+OPENAI_IMAGENET_TEMPLATES = (
+    lambda c: f'a bad photo of a {c}.',
+    lambda c: f'a photo of many {c}.',
+    lambda c: f'a sculpture of a {c}.',
+    lambda c: f'a photo of the hard to see {c}.',
+    lambda c: f'a low resolution photo of the {c}.',
+    lambda c: f'a rendering of a {c}.',
+    lambda c: f'graffiti of a {c}.',
+    lambda c: f'a bad photo of the {c}.',
+    lambda c: f'a cropped photo of the {c}.',
+    lambda c: f'a tattoo of a {c}.',
+    lambda c: f'the embroidered {c}.',
+    lambda c: f'a photo of a hard to see {c}.',
+    lambda c: f'a bright photo of a {c}.',
+    lambda c: f'a photo of a clean {c}.',
+    lambda c: f'a photo of a dirty {c}.',
+    lambda c: f'a dark photo of the {c}.',
+    lambda c: f'a drawing of a {c}.',
+    lambda c: f'a photo of my {c}.',
+    lambda c: f'the plastic {c}.',
+    lambda c: f'a photo of the cool {c}.',
+    lambda c: f'a close-up photo of a {c}.',
+    lambda c: f'a black and white photo of the {c}.',
+    lambda c: f'a painting of the {c}.',
+    lambda c: f'a painting of a {c}.',
+    lambda c: f'a pixelated photo of the {c}.',
+    lambda c: f'a sculpture of the {c}.',
+    lambda c: f'a bright photo of the {c}.',
+    lambda c: f'a cropped photo of a {c}.',
+    lambda c: f'a plastic {c}.',
+    lambda c: f'a photo of the dirty {c}.',
+    lambda c: f'a jpeg corrupted photo of a {c}.',
+    lambda c: f'a blurry photo of the {c}.',
+    lambda c: f'a photo of the {c}.',
+    lambda c: f'a good photo of the {c}.',
+    lambda c: f'a rendering of the {c}.',
+    lambda c: f'a {c} in a video game.',
+    lambda c: f'a photo of one {c}.',
+    lambda c: f'a doodle of a {c}.',
+    lambda c: f'a close-up photo of the {c}.',
+    lambda c: f'a photo of a {c}.',
+    lambda c: f'the origami {c}.',
+    lambda c: f'the {c} in a video game.',
+    lambda c: f'a sketch of a {c}.',
+    lambda c: f'a doodle of the {c}.',
+    lambda c: f'a origami {c}.',
+    lambda c: f'a low resolution photo of a {c}.',
+    lambda c: f'the toy {c}.',
+    lambda c: f'a rendition of the {c}.',
+    lambda c: f'a photo of the clean {c}.',
+    lambda c: f'a photo of a large {c}.',
+    lambda c: f'a rendition of a {c}.',
+    lambda c: f'a photo of a nice {c}.',
+    lambda c: f'a photo of a weird {c}.',
+    lambda c: f'a blurry photo of a {c}.',
+    lambda c: f'a cartoon {c}.',
+    lambda c: f'art of a {c}.',
+    lambda c: f'a sketch of the {c}.',
+    lambda c: f'a embroidered {c}.',
+    lambda c: f'a pixelated photo of a {c}.',
+    lambda c: f'itap of the {c}.',
+    lambda c: f'a jpeg corrupted photo of the {c}.',
+    lambda c: f'a good photo of a {c}.',
+    lambda c: f'a plushie {c}.',
+    lambda c: f'a photo of the nice {c}.',
+    lambda c: f'a photo of the small {c}.',
+    lambda c: f'a photo of the weird {c}.',
+    lambda c: f'the cartoon {c}.',
+    lambda c: f'art of the {c}.',
+    lambda c: f'a drawing of the {c}.',
+    lambda c: f'a photo of the large {c}.',
+    lambda c: f'a black and white photo of a {c}.',
+    lambda c: f'the plushie {c}.',
+    lambda c: f'a dark photo of a {c}.',
+    lambda c: f'itap of a {c}.',
+    lambda c: f'graffiti of the {c}.',
+    lambda c: f'a toy {c}.',
+    lambda c: f'itap of my {c}.',
+    lambda c: f'a photo of a cool {c}.',
+    lambda c: f'a photo of a small {c}.',
+    lambda c: f'a tattoo of the {c}.',
+)
+

preprocess.py

@@ -0,0 +1,42 @@
+## Imports
+import numpy as np
+import torch
+from PIL import Image
+import os.path
+import argparse
+from pathlib import Path
+import cv2
+from torch.nn import functional as F
+from torch.utils.data import DataLoader
+import tqdm
+import einops
+import plotly.express as px
+import torch.nn.functional as F
+import tqdm
+import json
+import albumentations
+import glob
+from torchvision import transforms
+
+
+def _convert_to_rgb(image):
+    return image.convert('RGB')
+
+def _resize(image):
+    image = np.array(image)
+    image = albumentations.augmentations.geometric.resize.LongestMaxSize(interpolation=Image.BICUBIC, 
+                                                                        max_size=224)(image=image)
+    return Image.fromarray(image['image'])
+
+preprocess = transforms.Compose([
+    _resize,
+    transforms.CenterCrop(size=(224, 224)),
+    _convert_to_rgb,
+])
+
+
+both_preprocess = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize(mean=(0.48145466, 0.4578275, 0.40821073), 
+                         std=(0.26862954, 0.26130258, 0.27577711)),
+])

preprocessor_config.json

@@ -0,0 +1,19 @@
+{
+  "crop_size": 336,
+  "do_center_crop": true,
+  "do_normalize": true,
+  "do_resize": true,
+  "feature_extractor_type": "CLIPFeatureExtractor",
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "resample": 3,
+  "size": 336
+}

preprocessor_config_bak.json

@@ -0,0 +1,16 @@
+{
+  "image_size": 224,
+  "image_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "image_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "interpolation": "bicubic",
+  "resize_mode": "center_crop"
+}
+

pretrained.py

@@ -0,0 +1,426 @@
+import hashlib
+import os
+import urllib
+import warnings
+from functools import partial
+from typing import Dict, Union
+
+from tqdm import tqdm
+
+
+try:
+    from huggingface_hub import hf_hub_download
+    hf_hub_download = partial(hf_hub_download, library_name="open_clip", library_version='2.20.0')
+    _has_hf_hub = True
+except ImportError:
+    hf_hub_download = None
+    _has_hf_hub = False
+
+
+def _pcfg(url='', hf_hub='', mean=None, std=None):
+    return dict(
+        url=url,
+        hf_hub=hf_hub,
+        mean=mean,
+        std=std,
+    )
+
+
+_RN50 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt"),
+    yfcc15m=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn50-quickgelu-yfcc15m-455df137.pt"),
+    cc12m=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn50-quickgelu-cc12m-f000538c.pt"),
+)
+
+_RN50_quickgelu = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt"),
+    yfcc15m=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn50-quickgelu-yfcc15m-455df137.pt"),
+    cc12m=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn50-quickgelu-cc12m-f000538c.pt"),
+)
+
+_RN101 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt"),
+    yfcc15m=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn101-quickgelu-yfcc15m-3e04b30e.pt"),
+)
+
+_RN101_quickgelu = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt"),
+    yfcc15m=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/rn101-quickgelu-yfcc15m-3e04b30e.pt"),
+)
+
+_RN50x4 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt"),
+)
+
+_RN50x16 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt"),
+)
+
+_RN50x64 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt"),
+)
+
+_VITB32 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt"),
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt"),
+    laion2b_e16=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-laion2b_e16-af8dbd0c.pth"),
+    laion2b_s34b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-laion2B-s34B-b79K/'),
+    # DataComp-M models
+    datacomp_m_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-DataComp.M-s128M-b4K/'),
+    commonpool_m_clip_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.clip-s128M-b4K/'),
+    commonpool_m_laion_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.laion-s128M-b4K/'),
+    commonpool_m_image_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.image-s128M-b4K/'),
+    commonpool_m_text_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.text-s128M-b4K/'),
+    commonpool_m_basic_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M.basic-s128M-b4K/'),
+    commonpool_m_s128m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.M-s128M-b4K/'),
+    # DataComp-S models
+    datacomp_s_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-DataComp.S-s13M-b4K/'),
+    commonpool_s_clip_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.clip-s13M-b4K/'),
+    commonpool_s_laion_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.laion-s13M-b4K/'),
+    commonpool_s_image_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.image-s13M-b4K/'),
+    commonpool_s_text_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.text-s13M-b4K/'),
+    commonpool_s_basic_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S.basic-s13M-b4K/'),
+    commonpool_s_s13m_b4k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-CommonPool.S-s13M-b4K/'),
+)
+
+_VITB32_quickgelu = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt"),
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e31-d867053b.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_32-quickgelu-laion400m_e32-46683a32.pt"),
+)
+
+_VITB16 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt"),
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e31-00efa78f.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16-laion400m_e32-55e67d44.pt"),
+    laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-laion2B-s34B-b88K/'),
+    # DataComp-L models
+    datacomp_l_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-DataComp.L-s1B-b8K/'),
+    commonpool_l_clip_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.clip-s1B-b8K/'),
+    commonpool_l_laion_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.laion-s1B-b8K/'),
+    commonpool_l_image_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.image-s1B-b8K/'),
+    commonpool_l_text_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.text-s1B-b8K/'),
+    commonpool_l_basic_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L.basic-s1B-b8K/'),
+    commonpool_l_s1b_b8k=_pcfg(hf_hub='laion/CLIP-ViT-B-16-CommonPool.L-s1B-b8K/'),
+)
+
+_VITB16_PLUS_240 = dict(
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e31-8fb26589.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_b_16_plus_240-laion400m_e32-699c4b84.pt"),
+)
+
+_VITL14 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt"),
+    laion400m_e31=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e31-69988bb6.pt"),
+    laion400m_e32=_pcfg(
+        "https://github.com/mlfoundations/open_clip/releases/download/v0.2-weights/vit_l_14-laion400m_e32-3d133497.pt"),
+    laion2b_s32b_b82k=_pcfg(
+        hf_hub='laion/CLIP-ViT-L-14-laion2B-s32B-b82K/',
+        mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
+    # DataComp-XL models
+    datacomp_xl_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-L-14-DataComp.XL-s13B-b90K/'),
+    commonpool_xl_clip_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-L-14-CommonPool.XL.clip-s13B-b90K/'),
+    commonpool_xl_laion_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-L-14-CommonPool.XL.laion-s13B-b90K/'),
+    commonpool_xl_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-L-14-CommonPool.XL-s13B-b90K/'),
+)
+
+_VITL14_336 = dict(
+    openai=_pcfg(
+        "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt"),
+)
+
+_VITH14 = dict(
+    laion2b_s32b_b79k=_pcfg(hf_hub='laion/CLIP-ViT-H-14-laion2B-s32B-b79K/'),
+)
+
+_VITg14 = dict(
+    laion2b_s12b_b42k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s12B-b42K/'),
+    laion2b_s34b_b88k=_pcfg(hf_hub='laion/CLIP-ViT-g-14-laion2B-s34B-b88K/'),
+)
+
+_VITbigG14 = dict(
+    laion2b_s39b_b160k=_pcfg(hf_hub='laion/CLIP-ViT-bigG-14-laion2B-39B-b160k/'),
+)
+
+_robertaViTB32 = dict(
+    laion2b_s12b_b32k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-roberta-base-laion2B-s12B-b32k/'),
+)
+
+_xlmRobertaBaseViTB32 = dict(
+    laion5b_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-xlm-roberta-base-laion5B-s13B-b90k/'),
+)
+
+_xlmRobertaLargeFrozenViTH14 = dict(
+    frozen_laion5b_s13b_b90k=_pcfg(hf_hub='laion/CLIP-ViT-H-14-frozen-xlm-roberta-large-laion5B-s13B-b90k/'),
+)
+
+_convnext_base = dict(
+    laion400m_s13b_b51k=_pcfg(hf_hub='laion/CLIP-convnext_base-laion400M-s13B-b51K/'),
+)
+
+_convnext_base_w = dict(
+    laion2b_s13b_b82k=_pcfg(hf_hub='laion/CLIP-convnext_base_w-laion2B-s13B-b82K/'),
+    laion2b_s13b_b82k_augreg=_pcfg(hf_hub='laion/CLIP-convnext_base_w-laion2B-s13B-b82K-augreg/'),
+    laion_aesthetic_s13b_b82k=_pcfg(hf_hub='laion/CLIP-convnext_base_w-laion_aesthetic-s13B-b82K/'),
+)
+
+_convnext_base_w_320 = dict(
+    laion_aesthetic_s13b_b82k=_pcfg(hf_hub='laion/CLIP-convnext_base_w_320-laion_aesthetic-s13B-b82K/'),
+    laion_aesthetic_s13b_b82k_augreg=_pcfg(hf_hub='laion/CLIP-convnext_base_w_320-laion_aesthetic-s13B-b82K-augreg/'),
+)
+
+_convnext_large_d = dict(
+    laion2b_s26b_b102k_augreg=_pcfg(hf_hub='laion/CLIP-convnext_large_d.laion2B-s26B-b102K-augreg/'),
+)
+
+_convnext_large_d_320 = dict(
+    laion2b_s29b_b131k_ft=_pcfg(hf_hub='laion/CLIP-convnext_large_d_320.laion2B-s29B-b131K-ft/'),
+    laion2b_s29b_b131k_ft_soup=_pcfg(hf_hub='laion/CLIP-convnext_large_d_320.laion2B-s29B-b131K-ft-soup/'),
+)
+
+_convnext_xxlarge = dict(
+    laion2b_s34b_b82k_augreg=_pcfg(hf_hub='laion/CLIP-convnext_xxlarge-laion2B-s34B-b82K-augreg/'),
+    laion2b_s34b_b82k_augreg_rewind=_pcfg(hf_hub='laion/CLIP-convnext_xxlarge-laion2B-s34B-b82K-augreg-rewind/'),
+    laion2b_s34b_b82k_augreg_soup=_pcfg(hf_hub='laion/CLIP-convnext_xxlarge-laion2B-s34B-b82K-augreg-soup/'),
+)
+
+_coca_VITB32 = dict(
+    laion2b_s13b_b90k=_pcfg(hf_hub='laion/CoCa-ViT-B-32-laion2B-s13B-b90k/'),
+    mscoco_finetuned_laion2b_s13b_b90k=_pcfg(hf_hub='laion/mscoco_finetuned_CoCa-ViT-B-32-laion2B-s13B-b90k/')
+)
+
+_coca_VITL14 = dict(
+    laion2b_s13b_b90k=_pcfg(hf_hub='laion/CoCa-ViT-L-14-laion2B-s13B-b90k/'),
+    mscoco_finetuned_laion2b_s13b_b90k=_pcfg(hf_hub='laion/mscoco_finetuned_CoCa-ViT-L-14-laion2B-s13B-b90k/')
+)
+
+
+_PRETRAINED = {
+    "RN50": _RN50,
+    "RN50-quickgelu": _RN50_quickgelu,
+    "RN101": _RN101,
+    "RN101-quickgelu": _RN101_quickgelu,
+    "RN50x4": _RN50x4,
+    "RN50x16": _RN50x16,
+    "RN50x64": _RN50x64,
+    "ViT-B-32": _VITB32,
+    "ViT-B-32-quickgelu": _VITB32_quickgelu,
+    "ViT-B-16": _VITB16,
+    "ViT-B-16-plus-240": _VITB16_PLUS_240,
+    "ViT-L-14": _VITL14,
+    "ViT-L-14-336": _VITL14_336,
+    "ViT-H-14": _VITH14,
+    "ViT-g-14": _VITg14,
+    "ViT-bigG-14": _VITbigG14,
+    "roberta-ViT-B-32": _robertaViTB32,
+    "xlm-roberta-base-ViT-B-32": _xlmRobertaBaseViTB32,
+    "xlm-roberta-large-ViT-H-14": _xlmRobertaLargeFrozenViTH14,
+    "convnext_base": _convnext_base,
+    "convnext_base_w": _convnext_base_w,
+    "convnext_base_w_320": _convnext_base_w_320,
+    "convnext_large_d": _convnext_large_d,
+    "convnext_large_d_320": _convnext_large_d_320,
+    "convnext_xxlarge": _convnext_xxlarge,
+    "coca_ViT-B-32": _coca_VITB32,
+    "coca_ViT-L-14": _coca_VITL14,
+    "EVA01-g-14": dict(
+        # from QuanSun/EVA-CLIP/EVA01_CLIP_g_14_psz14_s11B.pt
+        laion400m_s11b_b41k=_pcfg(hf_hub='timm/eva_giant_patch14_clip_224.laion400m_s11b_b41k/'),
+    ),
+    "EVA01-g-14-plus": dict(
+        # from QuanSun/EVA-CLIP/EVA01_CLIP_g_14_plus_psz14_s11B.pt
+        merged2b_s11b_b114k=_pcfg(hf_hub='timm/eva_giant_patch14_plus_clip_224.merged2b_s11b_b114k/'),
+    ),
+    "EVA02-B-16": dict(
+        # from QuanSun/EVA-CLIP/EVA02_CLIP_B_psz16_s8B.pt
+        merged2b_s8b_b131k=_pcfg(hf_hub='timm/eva02_base_patch16_clip_224.merged2b_s8b_b131k/'),
+    ),
+    "EVA02-L-14": dict(
+        # from QuanSun/EVA-CLIP/EVA02_CLIP_L_psz14_s4B.pt
+        merged2b_s4b_b131k=_pcfg(hf_hub='timm/eva02_large_patch14_clip_224.merged2b_s4b_b131k/'),
+    ),
+    "EVA02-L-14-336": dict(
+        # from QuanSun/EVA-CLIP/EVA02_CLIP_L_336_psz14_s6B.pt
+        merged2b_s6b_b61k=_pcfg(hf_hub='timm/eva02_large_patch14_clip_336.merged2b_s6b_b61k/'),
+    ),
+    "EVA02-E-14": dict(
+        # from QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_s4B.pt
+        laion2b_s4b_b115k=_pcfg(hf_hub='timm/eva02_enormous_patch14_clip_224.laion2b_s4b_b115k/'),
+    ),
+    "EVA02-E-14-plus": dict(
+        # from QuanSun/EVA-CLIP/EVA02_CLIP_E_psz14_plus_s9B.pt
+        laion2b_s9b_b144k=_pcfg(hf_hub='timm/eva02_enormous_patch14_plus_clip_224.laion2b_s9b_b144k/'),
+    )
+}
+
+
+def _clean_tag(tag: str):
+    # normalize pretrained tags
+    return tag.lower().replace('-', '_')
+
+
+def list_pretrained(as_str: bool = False):
+    """ returns list of pretrained models
+    Returns a tuple (model_name, pretrain_tag) by default or 'name:tag' if as_str == True
+    """
+    return [':'.join([k, t]) if as_str else (k, t) for k in _PRETRAINED.keys() for t in _PRETRAINED[k].keys()]
+
+
+def list_pretrained_models_by_tag(tag: str):
+    """ return all models having the specified pretrain tag """
+    models = []
+    tag = _clean_tag(tag)
+    for k in _PRETRAINED.keys():
+        if tag in _PRETRAINED[k]:
+            models.append(k)
+    return models
+
+
+def list_pretrained_tags_by_model(model: str):
+    """ return all pretrain tags for the specified model architecture """
+    tags = []
+    if model in _PRETRAINED:
+        tags.extend(_PRETRAINED[model].keys())
+    return tags
+
+
+def is_pretrained_cfg(model: str, tag: str):
+    if model not in _PRETRAINED:
+        return False
+    return _clean_tag(tag) in _PRETRAINED[model]
+
+
+def get_pretrained_cfg(model: str, tag: str):
+    if model not in _PRETRAINED:
+        return {}
+    model_pretrained = _PRETRAINED[model]
+    return model_pretrained.get(_clean_tag(tag), {})
+
+
+def get_pretrained_url(model: str, tag: str):
+    cfg = get_pretrained_cfg(model, _clean_tag(tag))
+    return cfg.get('url', '')
+
+
+def download_pretrained_from_url(
+        url: str,
+        cache_dir: Union[str, None] = None,
+):
+    if not cache_dir:
+        cache_dir = os.path.expanduser("~/.cache/clip")
+    os.makedirs(cache_dir, exist_ok=True)
+    filename = os.path.basename(url)
+
+    if 'openaipublic' in url:
+        expected_sha256 = url.split("/")[-2]
+    elif 'mlfoundations' in url:
+        expected_sha256 = os.path.splitext(filename)[0].split("-")[-1]
+    else:
+        expected_sha256 = ''
+
+    download_target = os.path.join(cache_dir, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if expected_sha256:
+            if hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
+                return download_target
+            else:
+                warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+        else:
+            return download_target
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.headers.get("Content-Length")), ncols=80, unit='iB', unit_scale=True) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if expected_sha256 and not hashlib.sha256(open(download_target, "rb").read()).hexdigest().startswith(expected_sha256):
+        raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def has_hf_hub(necessary=False):
+    if not _has_hf_hub and necessary:
+        # if no HF Hub module installed, and it is necessary to continue, raise error
+        raise RuntimeError(
+            'Hugging Face hub model specified but package not installed. Run `pip install huggingface_hub`.')
+    return _has_hf_hub
+
+
+def download_pretrained_from_hf(
+        model_id: str,
+        filename: str = 'open_clip_pytorch_model.bin',
+        revision=None,
+        cache_dir: Union[str, None] = None,
+):
+    has_hf_hub(True)
+    cached_file = hf_hub_download(model_id, filename, revision=revision, cache_dir=cache_dir)
+    return cached_file
+
+
+def download_pretrained(
+        cfg: Dict,
+        force_hf_hub: bool = False,
+        cache_dir: Union[str, None] = None,
+):
+    target = ''
+    if not cfg:
+        return target
+
+    download_url = cfg.get('url', '')
+    download_hf_hub = cfg.get('hf_hub', '')
+    if download_hf_hub and force_hf_hub:
+        # use HF hub even if url exists
+        download_url = ''
+
+    if download_url:
+        target = download_pretrained_from_url(download_url, cache_dir=cache_dir)
+    elif download_hf_hub:
+        has_hf_hub(True)
+        # we assume the hf_hub entries in pretrained config combine model_id + filename in
+        # 'org/model_name/filename.pt' form. To specify just the model id w/o filename and
+        # use 'open_clip_pytorch_model.bin' default, there must be a trailing slash 'org/model_name/'.
+        model_id, filename = os.path.split(download_hf_hub)
+        if filename:
+            target = download_pretrained_from_hf(model_id, filename=filename, cache_dir=cache_dir)
+        else:
+            target = download_pretrained_from_hf(model_id, cache_dir=cache_dir)
+
+    return target

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c7a812f61be88b4148e6c910ea245178ff3663263d54680cdb99dd6bcaed9b32
+size 1711950230

requirements.txt

@@ -0,0 +1,5 @@
+torch>=1.9.0
+transformers>=4.21.0
+torchvision>=0.10.0
+Pillow
+numpy

shared.py

@@ -0,0 +1,616 @@
+import os
+import subprocess
+import re
+import random
+import matplotlib.pyplot as plt
+import json
+def get_gpu_memory_usage():
+    """Returns a list of GPU memory usage in MB."""
+    try:
+        # Run nvidia-smi command and capture the output
+        result = subprocess.run(['nvidia-smi', '--query-gpu=memory.used', '--format=csv,nounits,noheader'],
+                                stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+
+        # Check if the command was successful
+        if result.returncode != 0:
+            raise RuntimeError(f"nvidia-smi command failed with error: {result.stderr}")
+
+        # Parse the output to get memory usage
+        memory_usages = [int(x) for x in result.stdout.strip().split('\n')]
+        return memory_usages
+    except Exception as e:
+        print(f"Error querying GPU memory usage: {e}")
+        return []
+
+def set_cuda_visible_device():
+    """Sets the CUDA_VISIBLE_DEVICES environment variable to the GPU with the smallest memory usage."""
+    memory_usages = get_gpu_memory_usage()
+
+    if not memory_usages:
+        print("No GPU memory usage data available.")
+        return
+
+    # Find the index of the GPU with the smallest memory usage
+    min_memory_index = memory_usages.index(min(memory_usages))
+
+    # Set the CUDA_VISIBLE_DEVICES environment variable
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(min_memory_index)
+    print(f"Set CUDA_VISIBLE_DEVICES to GPU {min_memory_index} with {memory_usages[min_memory_index]} MB used.")
+
+    return str(min_memory_index)
+
+os.environ["ASN_ROOT_DIR"] = "/home/nickj/asn/second_order_lens"
+os.chdir(os.environ["ASN_ROOT_DIR"])
+
+import numpy as np
+import torch
+from PIL import Image
+import os.path
+import argparse
+from pathlib import Path
+
+from tqdm import tqdm
+from utils.factory import create_model_and_transforms, get_tokenizer
+from PIL import Image, ImageDraw
+
+def get_model(model_name = "ViT-B/16", pretrained = "openai", device = "cuda:0"):
+  torch.multiprocessing.set_sharing_strategy("file_system")
+  model, _, preprocess = create_model_and_transforms(
+      model_name, pretrained=pretrained, force_quick_gelu=True,
+  )
+  model.to(device)
+  model.eval()
+  context_length = model.context_length
+  vocab_size = model.vocab_size
+
+  return {
+    "model": model,
+    "model_name": model_name,
+    "pretrained": pretrained,
+    "preprocess": preprocess,
+    "context_length": context_length,
+    "vocab_size": vocab_size
+  }
+
+img_path = "/datasets/ilsvrc_2024-01-04_1913/val/n04398044/ILSVRC2012_val_00042447.JPEG"
+# img_path = "./sample.jpeg"
+def load_images(preprocess, image_folder = "/datasets/ilsvrc/current/val", count = 100, images_only = True):
+  file_list = []
+
+  for root, dirs, files in os.walk(image_folder):
+      for file in files:
+          file_list.append(os.path.join(root, file))
+
+  if count > len(file_list):
+    sampled_files = file_list
+  else:
+    sampled_files = random.sample(file_list, count)
+
+  image_files = []
+
+  for filename in sampled_files:
+    image_files.append(preprocess(Image.open(filename)))
+  if images_only:
+    return image_files
+  else:
+    return image_files, sampled_files
+
+def calc_neuron_potentials(model, attn_layers = (1, 2), include_layernorm = True):
+  # Calculates the attention-shifting potential scores for every neuron to the attention heads defined by the given layers (relative to the MLP layer)
+
+  embed_dim = model.visual.transformer.resblocks[0].attn.out_proj.in_features
+  num_heads = model.visual.transformer.resblocks[0].attn.num_heads
+  head_dim = embed_dim // num_heads
+  layers = len(model.visual.transformer.resblocks)
+
+  results = dict()
+
+  for neuron_layer in tqdm(range(layers), desc = "Calculating attention shifting potentials"):
+    neuron_projection = model.visual.transformer.resblocks[neuron_layer].state_dict()["mlp.c_proj.weight"]
+    for l_attn in range(min(layers, neuron_layer + attn_layers[0]), min(layers, neuron_layer + attn_layers[1])):
+      ln_vector = model.visual.transformer.resblocks[l_attn].ln_1.state_dict()["weight"]
+      attn_matrix = model.visual.transformer.resblocks[l_attn].state_dict()["attn.in_proj_weight"]
+      W_Q, W_K, W_V = (attn_matrix[:embed_dim].reshape(num_heads, head_dim, -1),
+                       attn_matrix[embed_dim:2*embed_dim].reshape(num_heads, head_dim, -1),
+                       attn_matrix[2*embed_dim:].reshape(num_heads, head_dim, -1))
+
+      for head_idx in range(num_heads):
+        W_Q_h, W_K_h = W_Q[head_idx], W_K[head_idx]
+        effects = []
+        for i in range(neuron_projection.shape[1]):
+          if include_layernorm:
+            neuron_attn_effect = torch.norm(W_Q_h.T @ W_K_h @ (neuron_projection[:, i] * ln_vector))
+          else:
+            neuron_attn_effect = torch.norm(W_Q_h.T @ W_K_h @ neuron_projection[:, i])
+          effects.append(neuron_attn_effect)
+
+        results[(neuron_layer, l_attn, head_idx)] = torch.tensor(effects)
+  return results
+
+def calc_top_asns(shift_potentials, top_k = 10, per = "layer", layers_away = 1):
+  num_layers = max([key[1] for key in shift_potentials.keys()]) + 1 # the last layer has no ASNs by definition
+  num_heads = max([key[2] for key in shift_potentials.keys()])
+
+  top_asns = []
+  for layer in range(num_layers - layers_away):
+    if per == "layer":
+      potentials = []
+      for head_idx in range(num_heads):
+        potentials.append(shift_potentials[(layer, layer + layers_away, head_idx)])
+      potentials = torch.max(torch.stack(potentials, dim = 0), dim = 0).values
+      _, sorted_indices = torch.sort(potentials, descending = True)
+      top_asns.append(sorted_indices[:top_k].tolist())
+    elif per == "head":
+      top_layer_asns = []
+      for head_idx in range(num_heads):
+        _, sorted_indices = torch.sort(shift_potentials[(layer, layer + layers_away, head_idx)], descending = True)
+        top_layer_asns.append(sorted_indices[:top_k].tolist())
+      top_asns.append(top_layer_asns)
+    else:
+      raise ValueError(f"Invalid per value: {per}")
+  return top_asns
+
+def aggregate_attn_map(attn_map, layer, head):
+  num_tokens = attn_map.shape[-1]
+  assert (num_tokens - 1) ** 0.5 % 1 == 0, "num_tokens - 1 is not a perfect square"
+
+  num_patches = int((num_tokens - 1) ** 0.5)
+  aggregate_scores = torch.sum(attn_map[:, layer, head, 1:, 1:], dim = 1).reshape((1, num_patches, num_patches))
+  return aggregate_scores
+
+def attn_map_cls_token(attn_map, layer, head):
+  # Gets the attention map for the CLS token
+  num_tokens = attn_map.shape[-1]
+  assert (num_tokens - 1) ** 0.5 % 1 == 0, "num_tokens - 1 is not a perfect square"
+
+  num_patches = int((num_tokens - 1) ** 0.5)
+  attn_map_reshaped = attn_map[:, layer, head, 0, 1:].reshape((1, num_patches, num_patches))
+  return attn_map_reshaped
+
+def visualize_attn_shift(attn_map1, attn_map2, image, display=True, out=None, min_diff=None, max_diff=None):
+  import matplotlib.pyplot as plt
+  import numpy as np
+
+  # Subtract attn_map1 from attn_map2
+  diff_map = attn_map2 - attn_map1
+
+  # Convert the image to RGBA
+  image = image.convert("RGBA")
+  overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
+  draw = ImageDraw.Draw(overlay)
+
+  # Calculate the size of each attention block
+  block_size_x = image.size[0] / diff_map.shape[0]
+  block_size_y = image.size[1] / diff_map.shape[1]
+
+  # Create a colormap
+  cmap = plt.get_cmap('coolwarm_r')  # 'cool' colormap for lighter to darker
+
+  # Get the min and max values for scaling the colormap
+  if max_diff is None:
+    max_diff = diff_map.max()
+  if min_diff is None:
+    min_diff = diff_map.min()
+
+  for i in range(diff_map.shape[0]):
+    for j in range(diff_map.shape[1]):
+      # Get the color from the colormap
+      intensity = diff_map[i, j]
+      normalized_intensity = (intensity - min_diff) / (max_diff - min_diff)  # Scale to [0, 1]
+      rgba_color = cmap(1 - normalized_intensity)  # Invert the normalized intensity
+      color = tuple(int(c * 255) for c in rgba_color[:3]) + (int(rgba_color[3] * 128),)
+
+      # Draw the rectangle on the overlay with transparency
+      draw.rectangle(
+        [j * block_size_x, i * block_size_y, (j + 1) * block_size_x, (i + 1) * block_size_y],
+        fill=color  # Add transparency to the color
+      )
+
+  # Composite the overlay with the original image
+  combined = Image.alpha_composite(image, overlay)
+
+  if display:
+    # Display the result
+    combined.show()
+
+    # Show the color scale
+    plt.figure(figsize=(6, 1))
+    plt.imshow([np.linspace(min_diff, max_diff, 256)], cmap='coolwarm_r', aspect='auto')
+    plt.gca().set_visible(False)
+    plt.colorbar(orientation="horizontal")
+    plt.show()
+
+  if out is not None:
+    combined.save(out)
+
+  return combined
+
+def visualize_attn_shift_binary(attn_map1, attn_map2, image, display=True, out=None):
+  # Creates a visualization of the attention shift where green = positive, red = negative values.
+  # This is useful when there are outliers in the difference map causing the middle values around 0 to be messed into one color
+  # Subtract attn_map1 from attn_map2
+  diff_map = attn_map2 - attn_map1
+
+  # Normalize the difference map to range [0, 1] for visualization
+  diff_map_normalized = (diff_map - diff_map.min()) / (diff_map.max() - diff_map.min())
+  # Convert the image to RGBA
+  image = image.convert("RGBA")
+  overlay = Image.new("RGBA", image.size, (0, 0, 0, 0))
+  draw = ImageDraw.Draw(overlay)
+
+  # Calculate the size of each attention block
+  block_size_x = image.size[0] / diff_map.shape[0]
+  block_size_y = image.size[1] / diff_map.shape[1]
+
+  for i in range(diff_map.shape[0]):
+    for j in range(diff_map.shape[1]):
+      # Calculate the color intensity based on the difference
+      intensity = diff_map_normalized[i, j]
+      alpha = int(255 * 0.5)  # Tone down the alpha to 50%
+      if diff_map[i, j] > 0:
+        color = (0, int(255 * intensity), 0, alpha)  # Green for positive
+      else:
+        color = (int(255 * (1 - intensity)), 0, 0, alpha)  # Red for negative
+
+      # Draw the rectangle on the overlay
+      draw.rectangle(
+        [j * block_size_x, i * block_size_y, (j + 1) * block_size_x, (i + 1) * block_size_y],
+        fill=color
+      )
+
+  # Composite the overlay with the original image
+  combined = Image.alpha_composite(image, overlay)
+
+  if display:
+    # Display the result
+    combined.show()
+
+  if out is not None:
+    combined.save(out)
+
+  return combined
+
+def is_outlier(mean, std, value):
+  return value < mean - 2 * std or value > mean + 2 * std
+
+
+def get_neuron_activations(images, prs_group, model, device = "cuda:0"):
+  # Returns neuron activations in shape (num_images, num_layers, num_patches, num_neurons)
+  random_neuron_acts = []
+  for image in tqdm(images, desc="Processing images"):
+    prs_group.reinit()
+    image_input = image.unsqueeze(0).to(device)
+    representation = model.encode_image(
+      image_input, attn_method="head", normalize=False
+    )
+    prs_group.finalize()
+    gelu_outs = prs_group.post_gelu_outputs()
+    random_neuron_acts.append(gelu_outs)
+  random_neuron_acts = torch.stack(random_neuron_acts, dim = 0)
+  return random_neuron_acts
+
+def normalize_array(arr):
+  min_val = np.min(arr)
+  max_val = np.max(arr)
+  # Avoid division by zero if all values are the same
+  if max_val - min_val == 0:
+      return np.zeros_like(arr)
+  normalized_arr = (arr - min_val) / (max_val - min_val)
+  return normalized_arr
+
+def np_l2(arr1, arr2):
+  return np.linalg.norm(arr1 - arr2)
+
+def best_class(classifier, representation):
+  cs = torch.cosine_similarity(classifier, representation.permute(1, 0), dim = 0)
+  return torch.argmax(cs).item(), cs[torch.argmax(cs).item()].item()
+
+def load_group_attn_shifts(timestamp):
+  # Load from Supp1B
+  results_dir = "./results/supp1B"
+  # dirs = [os.path.join(results_dir, d) for d in os.listdir(results_dir)
+  #         if os.path.isdir(os.path.join(results_dir, d))]
+  # latest_dir = max(dirs, key=os.path.getmtime)
+
+  latest_dir = os.path.join(results_dir, timestamp)
+  print(f"Using latest results directory: {latest_dir}")
+
+  # Load metadata
+  with open(os.path.join(latest_dir, "metadata.json"), "r") as f:
+      metadata = json.load(f)
+
+  # Load memory-mapped files
+  attn_maps = np.memmap(os.path.join(latest_dir, "attention_maps.mmap"),
+                      dtype=np.float32,
+                      mode='r',
+                      shape=tuple(metadata["attention_maps_shape"]))
+
+  resblocks = np.memmap(os.path.join(latest_dir, "resblocks.mmap"),
+                      dtype=np.float32,
+                      mode='r',
+                      shape=tuple(metadata["resblocks_shape"]))
+
+  # Get file list from metadata
+  file_list = metadata.get("file_list", [])
+
+  # Get top_k values
+  top_k_values = metadata.get("top_k_values", [0])
+
+  return {
+    "attn_maps": attn_maps,
+    "resblocks": resblocks,
+    "metadata": metadata,
+    "file_list": file_list,
+    "top_k_values": top_k_values,
+    "num_layers": metadata.get("num_layers", 0),
+    "num_images": metadata.get("num_images", 0),
+    "num_heads": metadata.get("num_heads", 0)
+  }
+
+def load_individual_attn_shifts(timestamp, supp = "supp1D"):
+  results_dir = f"./results/{supp}"
+  # dirs = [os.path.join(results_dir, d) for d in os.listdir(results_dir)
+  #         if os.path.isdir(os.path.join(results_dir, d))]
+  # latest_dir = max(dirs, key=os.path.getmtime)
+
+  latest_dir = os.path.join(results_dir, timestamp)
+  print(f"Using latest results directory: {latest_dir}")
+
+  # Load metadata
+  with open(os.path.join(latest_dir, "metadata.json"), "r") as f:
+      metadata = json.load(f)
+
+  # Load memory-mapped files
+  attn_maps = np.memmap(os.path.join(latest_dir, "attention_maps.mmap"),
+                      dtype=np.float32,
+                      mode='r',
+                      shape=tuple(metadata["attention_maps_shape"]))
+
+  baseline_attn_maps = np.memmap(os.path.join(latest_dir, "baseline_attention_maps.mmap"),
+                      dtype=np.float32,
+                      mode='r',
+                      shape=tuple(metadata["baseline_attention_maps_shape"]))
+
+  neuron_activations = np.memmap(os.path.join(latest_dir, "neuron_activations.mmap"),
+                      dtype=np.float32,
+                      mode='r',
+                      shape=tuple(metadata["neuron_activations_shape"]))
+
+  baseline_neuron_activations = np.memmap(os.path.join(latest_dir, "baseline_neuron_activations.mmap"),
+                      dtype=np.float32,
+                      mode='r',
+                      shape=tuple(metadata["baseline_neuron_activations_shape"]))
+
+  ablated_neurons = np.memmap(os.path.join(latest_dir, "ablated_neurons.mmap"),
+                      dtype=np.float32,
+                      mode='r',
+                      shape=tuple(metadata["ablated_neurons_shape"]))
+
+  # Get file list from metadata
+  file_list = metadata.get("file_list", [])
+
+  # Get k value
+  k = metadata.get("k", 25)
+
+  return {
+    "attn_maps": attn_maps,
+    "baseline_attn_maps": baseline_attn_maps,
+    "neuron_activations": neuron_activations,
+    "baseline_neuron_activations": baseline_neuron_activations,
+    "ablated_neurons": ablated_neurons,
+    "metadata": metadata,
+    "file_list": file_list,
+    "k": k,
+    "num_layers": metadata.get("num_layers", 12),
+    "num_images": metadata.get("num_images", 100),
+    "model_name": metadata.get("model_name", "ViT-B-16"),
+    "pretrained": metadata.get("pretrained", "openai")
+  }
+
+def find_register_neurons_cuda(model, preprocess, prs_group, register_norm_threshold = 30, highest_layer = -1, device = "cuda:0", processed_image_cnt = 500):
+    num_layers = len(model.visual.transformer.resblocks)
+    highest_layer = num_layers - 1 if highest_layer == -1 else highest_layer
+    num_neurons = model.visual.transformer.resblocks[0].mlp.state_dict()["c_proj.weight"].shape[1]
+    random_images = load_images(preprocess, count=processed_image_cnt)
+    neuron_scores = torch.zeros((len(random_images), num_layers, num_neurons), device=device)
+    alignment_scores = torch.zeros((len(random_images), num_layers, num_neurons), device=device)
+    image_count = 0
+
+    for i in tqdm(range(len(random_images)), desc="Processing random images"):
+        image = random_images[i].unsqueeze(0).to(device)
+        prs_group.reinit()
+
+        with torch.inference_mode():
+            representation = model.encode_image(
+                image, attn_method="head", normalize=False
+            )
+        prs_group.finalize()
+
+        baseline_neuron_acts = prs_group.post_gelu_outputs().to(device)
+        baseline_resblock_outputs = prs_group.resblock_outputs().to(device)
+
+        # Calculate norm map using torch
+        norm_map = torch.norm(baseline_resblock_outputs[-1], dim=1)
+        filtered_norms = norm_map.clone()
+        filtered_norms[filtered_norms < register_norm_threshold] = 0
+
+        # Get register locations as a tensor
+        register_locations = torch.where(filtered_norms > register_norm_threshold)[0]
+
+        if len(register_locations) == 0:
+            continue
+
+        image_count += 1
+
+        # Process all layers vectorized
+        for layer in range(num_layers):
+            # Get absolute activations for all neurons in this layer
+            act_layer = torch.abs(baseline_neuron_acts[layer])  # Shape: [seq_len, num_neurons]
+
+            # Check sparsity condition for all neurons at once
+            sparse_neurons = torch.sum(act_layer < 0.5, dim=0) >= 0.5 * act_layer.shape[0]  # Shape: [num_neurons]
+
+            # Skip computation if no neurons meet the condition
+            if not torch.any(sparse_neurons):
+                continue
+
+            # Get values at register locations for all neurons simultaneously
+            # This creates a tensor of shape [num_register_locations, num_neurons]
+            register_values = act_layer[register_locations]
+
+            # For neurons that pass sparsity condition, compute mean at register locations
+            # First, compute mean for all neurons (this is fast)
+            neuron_means = register_values.mean(dim=0)  # Shape: [num_neurons]
+
+            # Then zero out means for neurons that don't pass sparsity condition
+            neuron_means = neuron_means * sparse_neurons.float()
+
+            # Store in score tensor
+            neuron_scores[i, layer] = neuron_means
+
+    # Rest of the code remains the same
+    mean_neuron_scores = neuron_scores[:image_count].mean(dim=0)
+    mean_alignment_scores = alignment_scores[:image_count].mean(dim=0)
+
+    # Flatten and find top values
+    flattened_scores = mean_neuron_scores.flatten()
+    sorted_values, sorted_indices = torch.sort(flattened_scores, descending=True)
+
+    flattened_alignment = mean_alignment_scores.flatten()
+    sorted_alignment_values, sorted_alignment_indices = torch.sort(flattened_alignment, descending=True)
+
+    # Convert indices to layer/neuron pairs
+    top_indices = [(idx.item() // num_neurons, idx.item() % num_neurons) for idx in sorted_indices]
+    top_alignment_indices = [(idx.item() // num_neurons, idx.item() % num_neurons) for idx in sorted_alignment_indices]
+
+    register_norms = [
+        (layer, neuron, sorted_values[i].item())
+        for i, (layer, neuron) in enumerate(top_indices)
+        if layer <= highest_layer
+    ]
+
+    best_alignment_scores = [
+        (layer, neuron, sorted_alignment_values[i].item())
+        for i, (layer, neuron) in enumerate(top_alignment_indices)
+        if layer <= highest_layer
+    ]
+
+    return register_norms, best_alignment_scores
+
+def find_register_neurons(model, preprocess, prs_group, register_norm_threshold = 30, highest_layer = -1, device = "cuda:0", processed_image_cnt = 500):
+  num_layers = len(model.visual.transformer.resblocks)
+  highest_layer = num_layers - 1 if highest_layer == -1 else highest_layer
+  num_neurons = model.visual.transformer.resblocks[0].mlp.state_dict()["c_proj.weight"].shape[1]
+
+  random_images = load_images(preprocess, count = processed_image_cnt)
+  neuron_scores = torch.zeros((len(random_images), num_layers, num_neurons))
+  for i in tqdm(range(len(random_images)), desc="Processing random images"):
+    image = random_images[i].unsqueeze(0).to(device)
+
+    prs_group.reinit()
+    with torch.no_grad():
+      representation = model.encode_image(
+        image, attn_method="head", normalize=False
+      )
+    prs_group.finalize()
+
+    # Gather neuron activations and resblock outputs
+    baseline_neuron_acts = prs_group.post_gelu_outputs().cpu().numpy()
+    baseline_resblock_outputs = prs_group.resblock_outputs().cpu().numpy()
+
+    # Calculate norms of the last resblock outputs. Only consider patches of the activation maps that correspond with registers
+    norms = np.linalg.norm(baseline_resblock_outputs[-1], axis=1)
+    norms[norms < register_norm_threshold] = 0
+    register_locations = np.where(norms > register_norm_threshold)[0]
+
+    # register_neurons = []
+    for layer in range(num_layers):
+      for neuron in range(num_neurons):
+        neuron_map = baseline_neuron_acts[layer, :, neuron]
+        mask = np.zeros_like(neuron_map, dtype=bool)
+        mask[register_locations] = True
+        neuron_map[~mask] = 0
+        if np.any(neuron_map < 0):
+          continue
+        # dist = np.linalg.norm(normalize_array(norms) - normalize_array(neuron_map))
+        # register_neurons.append((layer, neuron, dist.item(), neuron_map[register_locations].mean()))
+
+        neuron_scores[i, layer, neuron] = torch.tensor(neuron_map[register_locations].mean())
+  mean_neuron_scores = neuron_scores.mean(dim=0)
+  # Flatten the 2D tensor to find global top values
+  flattened_scores = mean_neuron_scores.flatten()
+  sorted_values, sorted_indices = torch.sort(flattened_scores, descending=True)
+
+  # Convert flat indices back to 2D coordinates (layer, neuron)
+  top_indices = [(idx.item() // num_neurons, idx.item() % num_neurons) for idx in sorted_indices]
+
+  return [(layer, neuron, sorted_values[i].item()) for i, (layer, neuron) in enumerate(top_indices) if layer <= highest_layer]
+
+
+def plot_attn_maps(attn_maps, image_idx):
+
+  num_layers, num_heads, patch_height, patch_width = attn_maps.shape
+  print(f"Shape of image_shifts: {attn_maps.shape}")
+
+  # Create a grid of plots for all layers and heads
+  fig, axes = plt.subplots(num_layers, num_heads, figsize=(2*num_heads, 2*num_layers))
+  fig.suptitle(f'Attention Shift Maps for Image #{image_idx}', fontsize=16)
+
+  # Import the correct module for make_axes_locatable
+  from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+  # Plot each layer-head combination
+  for layer in range(num_layers):
+      # Determine min and max for this layer for consistent colorbar scaling within the layer
+      layer_vmin = attn_maps[layer].min().item()
+      layer_vmax = attn_maps[layer].max().item()
+
+      for head in range(num_heads):
+          # Get the current axis (handle both 2D and 1D cases)
+          if num_layers == 1 and num_heads == 1:
+              ax = axes
+          elif num_layers == 1:
+              ax = axes[head]
+          elif num_heads == 1:
+              ax = axes[layer]
+          else:
+              ax = axes[layer, head]
+
+          # Plot the attention shift map with layer-specific normalization
+          im = ax.imshow(attn_maps[layer, head], cmap='viridis', vmin=layer_vmin, vmax=layer_vmax)
+
+          # Remove ticks for cleaner appearance
+          ax.set_xticks([])
+          ax.set_yticks([])
+
+          # Add layer and head labels only on the edges
+          if head == 0:
+              ax.set_ylabel(f'Layer {layer}')
+          if layer == num_layers-1:
+              ax.set_xlabel(f'Head {head}')
+
+          # Add a colorbar for each layer (only once per row)
+          if head == num_heads-1:
+              # Create a colorbar that's properly sized relative to the plot
+              divider = make_axes_locatable(ax)
+              cax = divider.append_axes("right", size="5%", pad=0.05)
+              plt.colorbar(im, cax=cax)
+
+  # Adjust layout to make room for the colorbars
+  plt.tight_layout()
+  return plt
+
+def calculate_iou(output, target):
+  intersection = output * (output == target)
+  area_inter = intersection.sum().item()
+  area_pred = output.sum().item()
+  area_target = target.sum().item()
+  union = area_pred + area_target - area_inter
+  iou = area_inter / union
+  return area_inter, union, iou
+
+def calculate_pixel_accuracy(output, target):
+  correct = output * (output == target)
+  correct = correct.sum().item()
+  total = target.sum().item()
+  return correct, total, correct / total

special_tokens_map.json

@@ -0,0 +1 @@
+{"bos_token": {"content": "<|startoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true}, "pad_token": "<|endoftext|>"}

timm_model.py

@@ -0,0 +1,149 @@
+""" timm model adapter
+
+Wraps timm (https://github.com/rwightman/pytorch-image-models) models for use as a vision tower in CLIP model.
+"""
+import logging
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+try:
+    import timm
+    from timm.models.layers import Mlp, to_2tuple
+    try:
+        # old timm imports < 0.8.1
+        from timm.models.layers.attention_pool2d import RotAttentionPool2d
+        from timm.models.layers.attention_pool2d import AttentionPool2d as AbsAttentionPool2d
+    except ImportError:
+        # new timm imports >= 0.8.1
+        from timm.layers import RotAttentionPool2d
+        from timm.layers import AttentionPool2d as AbsAttentionPool2d
+except ImportError:
+    timm = None
+
+from misc import freeze_batch_norm_2d
+
+
+class TimmModel(nn.Module):
+    """ timm model adapter
+    """
+
+    def __init__(
+            self,
+            model_name,
+            embed_dim,
+            image_size=224,
+            pool='avg',
+            proj='linear',
+            proj_bias=False,
+            drop=0.,
+            drop_path=None,
+            patch_drop=None,
+            pretrained=False,
+    ):
+        super().__init__()
+        if timm is None:
+            raise RuntimeError("Please `pip install timm` to use timm models.")
+        self.image_size = to_2tuple(image_size)
+
+        # setup kwargs that may not be common across all models
+        timm_kwargs = {}
+        if drop_path is not None:
+            timm_kwargs['drop_path_rate'] = drop_path
+        if patch_drop is not None:
+            timm_kwargs['patch_drop_rate'] = patch_drop
+
+        custom_pool = pool in ('abs_attn', 'rot_attn')
+        if not proj and not custom_pool:
+            # use network classifier head as projection if no proj specified and no custom pooling used
+            self.trunk = timm.create_model(
+                model_name,
+                num_classes=embed_dim,
+                global_pool=pool,
+                pretrained=pretrained,
+                **timm_kwargs,
+            )
+            prev_chs = embed_dim
+        else:
+            self.trunk = timm.create_model(
+                model_name,
+                pretrained=pretrained,
+                **timm_kwargs,
+            )
+            feat_size = self.trunk.default_cfg.get('pool_size', None)
+            feature_ndim = 1 if not feat_size else 2
+            if custom_pool:
+                assert feature_ndim == 2
+                # if attn pooling used, remove both classifier and default pool
+                self.trunk.reset_classifier(0, global_pool='')
+            else:
+                # reset global pool if pool config set, otherwise leave as network default
+                reset_kwargs = dict(global_pool=pool) if pool else {}
+                self.trunk.reset_classifier(0, **reset_kwargs)
+            prev_chs = self.trunk.num_features
+
+        head_layers = OrderedDict()
+
+        # Add custom pooling to head
+        if pool == 'abs_attn':
+            head_layers['pool'] = AbsAttentionPool2d(prev_chs, feat_size=feat_size, out_features=embed_dim)
+            prev_chs = embed_dim
+        elif pool == 'rot_attn':
+            head_layers['pool'] = RotAttentionPool2d(prev_chs, out_features=embed_dim)
+            prev_chs = embed_dim
+
+        # NOTE attention pool ends with a projection layer, so proj should usually be set to '' if such pooling is used
+        if proj == 'linear':
+            head_layers['drop'] = nn.Dropout(drop)
+            head_layers['proj'] = nn.Linear(prev_chs, embed_dim, bias=proj_bias)
+        elif proj == 'mlp':
+            head_layers['mlp'] = Mlp(prev_chs, 2 * embed_dim, embed_dim, drop=(drop, 0), bias=(True, proj_bias))
+        else:
+            assert not proj, f'Unknown projection type {proj}.'
+
+        self.head = nn.Sequential(head_layers)
+
+    def lock(self, unlocked_groups=0, freeze_bn_stats=False):
+        """ lock modules
+        Args:
+            unlocked_groups (int): leave last n layer groups unlocked (default: 0)
+        """
+        if not unlocked_groups:
+            # lock full model
+            for param in self.trunk.parameters():
+                param.requires_grad = False
+            if freeze_bn_stats:
+                freeze_batch_norm_2d(self.trunk)
+        else:
+            # NOTE: partial freeze requires latest timm (master) branch and is subject to change
+            try:
+                # FIXME import here until API stable and in an official release
+                from timm.models.helpers import group_parameters, group_modules
+            except ImportError:
+                raise RuntimeError(
+                    'Please install latest timm `pip install git+https://github.com/rwightman/pytorch-image-models`')
+            matcher = self.trunk.group_matcher()
+            gparams = group_parameters(self.trunk, matcher)
+            max_layer_id = max(gparams.keys())
+            max_layer_id = max_layer_id - unlocked_groups
+            for group_idx in range(max_layer_id + 1):
+                group = gparams[group_idx]
+                for param in group:
+                    self.trunk.get_parameter(param).requires_grad = False
+            if freeze_bn_stats:
+                gmodules = group_modules(self.trunk, matcher, reverse=True)
+                gmodules = {k for k, v in gmodules.items() if v <= max_layer_id}
+                freeze_batch_norm_2d(self.trunk, gmodules)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        try:
+            self.trunk.set_grad_checkpointing(enable)
+        except Exception as e:
+            logging.warning('grad checkpointing not supported for this timm image tower, continuing without...')
+
+    def forward(self, x):
+        x = self.trunk(x)
+        x = self.head(x)
+        return x

tokenizer.py

@@ -0,0 +1,214 @@
+""" CLIP tokenizer
+
+Copied from https://github.com/openai/CLIP. Originally MIT License, Copyright (c) 2021 OpenAI.
+"""
+import gzip
+import html
+import os
+from functools import lru_cache
+from typing import Union, List
+
+import ftfy
+import regex as re
+import torch
+
+# https://stackoverflow.com/q/62691279
+import os
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "vocab/bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a significant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe(), special_tokens=None):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        if not special_tokens:
+            special_tokens = ['<start_of_text>', '<end_of_text>']
+        else:
+            special_tokens = ['<start_of_text>', '<end_of_text>'] + special_tokens
+        vocab.extend(special_tokens)
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {t:t for t in special_tokens}
+        special = "|".join(special_tokens)
+        self.pat = re.compile(special + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+        self.vocab_size = len(self.encoder)
+        self.all_special_ids = [self.encoder[t] for t in special_tokens]
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text
+
+
+_tokenizer = SimpleTokenizer()
+
+def decode(output_ids: torch.Tensor):
+    output_ids = output_ids.cpu().numpy()
+    return _tokenizer.decode(output_ids)
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77) -> torch.LongTensor:
+    """
+    Returns the tokenized representation of given input string(s)
+
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<start_of_text>"]
+    eot_token = _tokenizer.encoder["<end_of_text>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            tokens = tokens[:context_length]  # Truncate
+            tokens[-1] = eot_token
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result
+
+
+class HFTokenizer:
+    """HuggingFace tokenizer wrapper"""
+
+    def __init__(self, tokenizer_name: str):
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+
+    def save_pretrained(self, dest):
+        self.tokenizer.save_pretrained(dest)
+
+    def __call__(self, texts: Union[str, List[str]], context_length: int = 77) -> torch.Tensor:
+        # same cleaning as for default tokenizer, except lowercasing
+        # adding lower (for case-sensitive tokenizers) will make it more robust but less sensitive to nuance
+        if isinstance(texts, str):
+            texts = [texts]
+        texts = [whitespace_clean(basic_clean(text)) for text in texts]
+        input_ids = self.tokenizer(
+            texts,
+            return_tensors='pt',
+            max_length=context_length,
+            padding='max_length',
+            truncation=True,
+        ).input_ids
+        return input_ids

tokenizer_config.json

@@ -0,0 +1,34 @@
+{
+    "unk_token": {
+        "content": "<|endoftext|>",
+        "single_word": false,
+        "lstrip": false,
+        "rstrip": false,
+        "normalized": true,
+        "__type": "AddedToken"
+    },
+    "bos_token": {
+        "content": "<|startoftext|>",
+        "single_word": false,
+        "lstrip": false,
+        "rstrip": false,
+        "normalized": true,
+        "__type": "AddedToken"
+    },
+    "eos_token": {
+        "content": "<|endoftext|>",
+        "single_word": false,
+        "lstrip": false,
+        "rstrip": false,
+        "normalized": true,
+        "__type": "AddedToken"
+    },
+    "pad_token": "<|endoftext|>",
+    "add_prefix_space": false,
+    "errors": "replace",
+    "do_lower_case": true,
+    "name_or_path": "openai/clip-vit-base-patch32",
+    "model_max_length": 77,
+    "special_tokens_map_file": "/home/suraj/.cache/huggingface/transformers/18a566598f286c9139f88160c99f84eec492a26bd22738fa9cb44d5b7e0a5c76.cce1206abbad28826f000510f22f354e53e66a97f7c23745a7dfe27609cc07f5",
+    "tokenizer_class": "CLIPTokenizer"
+}

tokenizer_config_bak.json

@@ -0,0 +1,10 @@
+{
+  "tokenizer_class": "SimpleTokenizer",
+  "vocab_size": 49408,
+  "context_length": 77,
+  "bpe_path": "vocab/bpe_simple_vocab_16e6.txt.gz",
+  "special_tokens": [
+    "<start_of_text>",
+    "<end_of_text>"
+  ]
+}

transform.py

@@ -0,0 +1,133 @@
+import warnings
+from dataclasses import dataclass, asdict
+from typing import Any, Dict, Optional, Sequence, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as F
+
+from torchvision.transforms import Normalize, Compose, RandomResizedCrop, InterpolationMode, ToTensor, Resize, \
+    CenterCrop
+
+from constants import OPENAI_DATASET_MEAN, OPENAI_DATASET_STD
+
+
+@dataclass
+class AugmentationCfg:
+    scale: Tuple[float, float] = (0.9, 1.0)
+    ratio: Optional[Tuple[float, float]] = None
+    color_jitter: Optional[Union[float, Tuple[float, float, float]]] = None
+    interpolation: Optional[str] = None
+    re_prob: Optional[float] = None
+    re_count: Optional[int] = None
+    use_timm: bool = False
+
+
+class ResizeMaxSize(nn.Module):
+
+    def __init__(self, max_size, interpolation=InterpolationMode.BICUBIC, fn='max', fill=0):
+        super().__init__()
+        if not isinstance(max_size, int):
+            raise TypeError(f"Size should be int. Got {type(max_size)}")
+        self.max_size = max_size
+        self.interpolation = interpolation
+        self.fn = min if fn == 'min' else min
+        self.fill = fill
+
+    def forward(self, img):
+        if isinstance(img, torch.Tensor):
+            height, width = img.shape[:2]
+        else:
+            width, height = img.size
+        scale = self.max_size / float(max(height, width))
+        if scale != 1.0:
+            new_size = tuple(round(dim * scale) for dim in (height, width))
+            img = F.resize(img, new_size, self.interpolation)
+            pad_h = self.max_size - new_size[0]
+            pad_w = self.max_size - new_size[1]
+            img = F.pad(img, padding=[pad_w//2, pad_h//2, pad_w - pad_w//2, pad_h - pad_h//2], fill=self.fill)
+        return img
+
+
+def _convert_to_rgb(image):
+    return image.convert('RGB')
+
+
+def image_transform(
+        image_size: int,
+        is_train: bool,
+        mean: Optional[Tuple[float, ...]] = None,
+        std: Optional[Tuple[float, ...]] = None,
+        resize_longest_max: bool = False,
+        fill_color: int = 0,
+        aug_cfg: Optional[Union[Dict[str, Any], AugmentationCfg]] = None,
+):
+    mean = mean or OPENAI_DATASET_MEAN
+    if not isinstance(mean, (list, tuple)):
+        mean = (mean,) * 3
+
+    std = std or OPENAI_DATASET_STD
+    if not isinstance(std, (list, tuple)):
+        std = (std,) * 3
+
+    if isinstance(image_size, (list, tuple)) and image_size[0] == image_size[1]:
+        # for square size, pass size as int so that Resize() uses aspect preserving shortest edge
+        image_size = image_size[0]
+
+    if isinstance(aug_cfg, dict):
+        aug_cfg = AugmentationCfg(**aug_cfg)
+    else:
+        aug_cfg = aug_cfg or AugmentationCfg()
+    normalize = Normalize(mean=mean, std=std)
+    if is_train:
+        aug_cfg_dict = {k: v for k, v in asdict(aug_cfg).items() if v is not None}
+        use_timm = aug_cfg_dict.pop('use_timm', False)
+        if use_timm:
+            from timm.data import create_transform  # timm can still be optional
+            if isinstance(image_size, (tuple, list)):
+                assert len(image_size) >= 2
+                input_size = (3,) + image_size[-2:]
+            else:
+                input_size = (3, image_size, image_size)
+            # by default, timm aug randomly alternates bicubic & bilinear for better robustness at inference time
+            aug_cfg_dict.setdefault('interpolation', 'random')
+            aug_cfg_dict.setdefault('color_jitter', None)  # disable by default
+            train_transform = create_transform(
+                input_size=input_size,
+                is_training=True,
+                hflip=0.,
+                mean=mean,
+                std=std,
+                re_mode='pixel',
+                **aug_cfg_dict,
+            )
+        else:
+            train_transform = Compose([
+                RandomResizedCrop(
+                    image_size,
+                    scale=aug_cfg_dict.pop('scale'),
+                    interpolation=InterpolationMode.BICUBIC,
+                ),
+                _convert_to_rgb,
+                ToTensor(),
+                normalize,
+            ])
+            if aug_cfg_dict:
+                warnings.warn(f'Unused augmentation cfg items, specify `use_timm` to use ({list(aug_cfg_dict.keys())}).')
+        return train_transform
+    else:
+        if resize_longest_max:
+            transforms = [
+                ResizeMaxSize(image_size, fill=fill_color)
+            ]
+        else:
+            transforms = [
+                Resize(image_size, interpolation=InterpolationMode.BICUBIC),
+                CenterCrop(image_size),
+            ]
+        transforms.extend([
+            _convert_to_rgb,
+            ToTensor(),
+            normalize,
+        ])
+        return Compose(transforms)

transformer.py

@@ -0,0 +1,872 @@
+from collections import OrderedDict
+import math
+from typing import Callable, Optional, Sequence, Tuple, Text
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+from torch.utils.checkpoint import checkpoint
+import numbers
+import einops
+import numpy as np
+from misc import to_2tuple
+
+
+class LayerNorm(nn.Module):
+    """Subclass torch's LayerNorm (with cast back to input dtype)."""
+
+    def __init__(
+        self,
+        normalized_shape,
+        eps: float = 1e-5,
+        elementwise_affine: bool = True,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        self.normalized_shape = tuple(normalized_shape)
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = torch.nn.Parameter(
+                torch.empty(self.normalized_shape)
+            )
+            self.bias = torch.nn.Parameter(
+                torch.empty(self.normalized_shape)
+            )
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        assert self.normalized_shape == x.shape[-len(self.normalized_shape):]
+        dims = [-(i + 1) for i in range(len(self.normalized_shape))]
+        mean = x.mean(dim=dims, keepdim=True)
+        mean_x2 = (x**2).mean(dim=dims, keepdim=True)
+        var = mean_x2 - mean**2
+        x_norm = (x - mean) / torch.sqrt(var + self.eps)
+        if self.elementwise_affine:
+            x_norm = self.weight * x_norm + self.bias
+        return x_norm.to(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_values=1e-5, inplace=False):
+        super().__init__()
+        self.inplace = inplace
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x):
+        raise ValueError("Not implemented")
+        return x.mul_(self.gamma) if self.inplace else x * self.gamma
+
+
+class PatchDropout(nn.Module):
+    """
+    https://arxiv.org/abs/2212.00794
+    """
+
+    def __init__(self, prob, exclude_first_token=True):
+        super().__init__()
+        assert 0 <= prob < 1.0
+        self.prob = prob
+        self.exclude_first_token = exclude_first_token
+
+    def forward(self, x):
+        if not self.training or self.prob == 0.0:
+            return x
+
+        if self.exclude_first_token:
+            cls_tokens, x = x[:, :1], x[:, 1:]
+        else:
+            cls_tokens = torch.jit.annotate(torch.Tensor, x[:, :1])
+
+        batch = x.size()[0]
+        num_tokens = x.size()[1]
+
+        batch_indices = torch.arange(batch)
+        batch_indices = batch_indices[..., None]
+
+        keep_prob = 1 - self.prob
+        num_patches_keep = max(1, int(num_tokens * keep_prob))
+
+        rand = torch.randn(batch, num_tokens)
+        patch_indices_keep = rand.topk(num_patches_keep, dim=-1).indices
+
+        x = x[batch_indices, patch_indices_keep]
+
+        if self.exclude_first_token:
+            x = torch.cat((cls_tokens, x), dim=1)
+
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=True,
+        scaled_cosine=False,
+        scale_heads=False,
+        logit_scale_max=math.log(1.0 / 0.01),
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.scaled_cosine = scaled_cosine
+        self.scale_heads = scale_heads
+        assert dim % num_heads == 0, "dim should be divisible by num_heads"
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim**-0.5
+        self.logit_scale_max = logit_scale_max
+
+        self.in_proj_weight = nn.Parameter(torch.randn((dim * 3, dim)) * self.scale)
+        if qkv_bias:
+            self.in_proj_bias = nn.Parameter(torch.zeros(dim * 3))
+        else:
+            self.in_proj_bias = None
+
+        if self.scaled_cosine:
+            self.logit_scale = nn.Parameter(
+                torch.log(10 * torch.ones((num_heads, 1, 1)))
+            )
+        else:
+            self.logit_scale = None
+        self.attn_drop = nn.Dropout(attn_drop)
+        if self.scale_heads:
+            self.head_scale = nn.Parameter(torch.ones((num_heads, 1, 1)))
+        else:
+            self.head_scale = None
+        self.out_proj = nn.Linear(dim, dim)
+        self.out_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):
+        L, N, C = x.shape
+        q, k, v = F.linear(x, self.in_proj_weight, self.in_proj_bias).chunk(3, dim=-1)
+        q = q.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+        k = k.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+        v = v.contiguous().view(L, N * self.num_heads, -1).transpose(0, 1)
+
+        if self.logit_scale is not None:
+            attn = torch.bmm(
+                F.normalize(q, dim=-1), F.normalize(k, dim=-1).transpose(-1, -2)
+            )
+            logit_scale = torch.clamp(self.logit_scale, max=self.logit_scale_max).exp()
+            attn = attn.view(N, self.num_heads, L, L) * logit_scale
+            attn = attn.view(-1, L, L)
+        else:
+            q = q * self.scale
+            attn = torch.bmm(q, k.transpose(-1, -2))
+
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                new_attn_mask = torch.zeros_like(attn_mask, dtype=q.dtype)
+                new_attn_mask.masked_fill_(attn_mask, float("-inf"))
+                attn_mask = new_attn_mask
+            attn += attn_mask
+
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = torch.bmm(attn, v)
+        if self.head_scale is not None:
+            x = x.view(N, self.num_heads, L, C) * self.head_scale
+            x = x.view(-1, L, C)
+        x = x.transpose(0, 1).reshape(L, N, C)
+        x = self.out_proj(x)
+        x = self.out_drop(x)
+        return x
+
+
+class AttentionalPooler(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        context_dim: int,
+        n_head: int = 8,
+        n_queries: int = 256,
+        norm_layer: Callable = LayerNorm,
+    ):
+        super().__init__()
+        self.query = nn.Parameter(torch.randn(n_queries, d_model))
+        self.attn = nn.MultiheadAttention(
+            d_model, n_head, kdim=context_dim, vdim=context_dim
+        )
+        self.ln_q = norm_layer(d_model)
+        self.ln_k = norm_layer(context_dim)
+
+    def forward(self, x: torch.Tensor):
+        x = self.ln_k(x).permute(1, 0, 2)  # NLD -> LND
+        N = x.shape[1]
+        q = self.ln_q(self.query)
+        out = self.attn(self._repeat(q, N), x, x, need_weights=False)[0]
+        return out.permute(1, 0, 2)  # LND -> NLD
+
+    def _repeat(self, query, N: int):
+        return query.unsqueeze(1).repeat(1, N, 1)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        mlp_width: int,
+        act_layer: Callable = nn.GELU,
+        layer_id: Optional[int] = None,
+    ):
+        super().__init__()
+        self.c_fc = nn.Linear(d_model, mlp_width)
+        self.gelu = act_layer()
+        self.c_proj = nn.Linear(mlp_width, d_model)
+        self.layer_id = layer_id
+        
+    def forward(self, x, neuron_dict=None, num_register_tokens=0):
+        x = self.c_fc(x)
+        
+        # If we have a dictionary of modifications and this layer is in it
+        if neuron_dict is not None and self.layer_id in neuron_dict and num_register_tokens>0:
+            neurons = neuron_dict[self.layer_id]
+            
+            # Apply GELU to all activations
+            x_after_gelu = self.gelu(x)
+
+            original_activations = x_after_gelu.clone()
+            # Create new activation map for specified neurons
+            new_activation_map = torch.zeros(
+                (x_after_gelu.shape[0], x_after_gelu.shape[1], len(neurons)),
+                device=x_after_gelu.device, 
+            ).to(x_after_gelu.dtype)
+
+            max_values = torch.max(original_activations[:, :, neurons], dim=1, keepdim=True).values
+            
+            new_activation_map[:, -num_register_tokens:, :] = max_values
+            new_activation_map[:,0,:] = x_after_gelu[:,0,neurons]
+
+            x_after_gelu[:,:,neurons] = new_activation_map
+            x = x_after_gelu
+        else:
+            x = self.gelu(x)
+        
+        x = self.c_proj(x)
+        return x
+
+# TODO 여기가 custom attetion이 아니라는 점에서 문제가 발생한 것으로 보인다.
+class MultiheadAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv=False,
+        add_zero_attn=False,
+        kdim=None,
+        vdim=None,
+        batch_first=False,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim
+        self.q_out = nn.Identity()
+        self.k_out = nn.Identity()
+        self.v_out = nn.Identity()
+        self.qkv_out = nn.Identity()
+        self.attn_map = nn.Identity()
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.batch_first = batch_first
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+        self.in_proj_weight = nn.Parameter(torch.empty((3 * embed_dim, embed_dim)))
+
+        if bias:
+            self.in_proj_bias = nn.Parameter(torch.empty(3 * embed_dim))
+        else:
+            self.register_parameter("in_proj_bias", None)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        if add_bias_kv:
+            self.bias_k = nn.Parameter(torch.empty((1, 1, embed_dim)))
+            self.bias_v = nn.Parameter(torch.empty((1, 1, embed_dim)))
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+    def forward_direct(self, x, attn_mask=None):
+        B, N, C = x.shape
+        qkv = x @ self.in_proj_weight.T + self.in_proj_bias
+        qkv = self.qkv_out(qkv)
+        qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        # B, S, 3, H, d -> 3, B, H, S, d  batch first computation
+        # 이 지점 때문에 연산 결과에 차이가 생기는 거 같은데? 
+        q, k, v = qkv.unbind(0)
+
+        q = self.q_out(q)
+        k = self.k_out(k)
+        v = self.v_out(v)
+
+        dk = q.size()[-1]
+        q = q / math.sqrt(dk)
+        attn = q @ k.transpose(-2, -1)
+        if attn_mask is not None:
+            attn += attn_mask
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_map(attn)
+        x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = x @ self.out_proj.weight.T + self.out_proj.bias
+        return x
+
+    def _split_qkv_weight(self):
+        q_weight, k_weight, v_weight = (
+            self.in_proj_weight[: self.embed_dim].reshape(
+                self.num_heads, self.head_dim, -1
+            ),
+            self.in_proj_weight[self.embed_dim : self.embed_dim * 2].reshape(
+                self.num_heads, self.head_dim, -1
+            ),
+            self.in_proj_weight[self.embed_dim * 2 :].reshape(
+                self.num_heads, self.head_dim, -1
+            ),
+        )
+        return q_weight, k_weight, v_weight
+
+    def _split_qkv_bias(self):
+        q_bias, k_bias, v_bias = (
+            self.in_proj_bias[: self.embed_dim].reshape(
+                1, self.num_heads, 1, self.head_dim
+            ),
+            self.in_proj_bias[self.embed_dim : self.embed_dim * 2].reshape(
+                1, self.num_heads, 1, self.head_dim
+            ),
+            self.in_proj_bias[self.embed_dim * 2 :].reshape(
+                1, self.num_heads, 1, self.head_dim
+            ),
+        )
+        return q_bias, k_bias, v_bias
+
+    def forward_qkv(self, x, attn_mask=None):
+        B, N, C = x.shape
+        q_weight, k_weight, v_weight = (
+            self.in_proj_weight[: self.embed_dim],
+            self.in_proj_weight[self.embed_dim : self.embed_dim * 2],
+            self.in_proj_weight[self.embed_dim * 2 :],
+        )
+        q_bias, k_bias, v_bias = (
+            self.in_proj_bias[: self.embed_dim],
+            self.in_proj_bias[self.embed_dim : self.embed_dim * 2],
+            self.in_proj_bias[self.embed_dim * 2 :],
+        )
+        q = (x @ q_weight.T + q_bias).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        k = (x @ k_weight.T + k_bias).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        v = (x @ v_weight.T + v_bias).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        
+        dk = q.size()[-1]
+        q = q / math.sqrt(dk)
+        attn = q @ k.transpose(-2, -1)
+        if attn_mask is not None:
+            attn += attn_mask
+        attn = attn.softmax(dim=-1)
+        x = torch.einsum("bhnm,bhmc->bhnmc", attn, v)
+        x = x.sum(axis=3).transpose(1, 2).reshape(B, N, C)
+        x = x @ self.out_proj.weight.T + self.out_proj.bias
+        return x
+
+    def forward_per_head(self, x, attn_mask=None):
+        B, N, C = x.shape
+        q_weight, k_weight, v_weight = self._split_qkv_weight()
+        q_bias, k_bias, v_bias = self._split_qkv_bias()
+        q = torch.einsum("bnc,hdc->bhnd", x, q_weight) + q_bias
+        k = torch.einsum("bnc,hdc->bhnd", x, k_weight) + k_bias
+        v = torch.einsum("bnc,hdc->bhnd", x, v_weight) + v_bias
+        
+        dk = q.size()[-1]
+        q = q / math.sqrt(dk)
+        attn = q @ k.transpose(-2, -1)
+        if attn_mask is not None:
+            attn += attn_mask
+        attn = attn.softmax(dim=-1)
+        x = torch.einsum("bhnm,bhmc->bnmhc", attn, v)
+        x = torch.einsum(
+            "bnmhc,dhc->bnmhd",
+            x,
+            self.out_proj.weight.reshape(self.embed_dim, self.num_heads, self.head_dim),
+        )
+        x = x.sum(axis=[2, 3]) + self.out_proj.bias
+        return x
+
+    def _get_ov_circuit(self):
+        reshaped_o = self.out_proj.weight.reshape(
+            self.embed_dim, self.num_heads, self.head_dim
+        )
+        _, _, v_weight = self._split_qkv_weight()
+        _, _, v_bias = self._split_qkv_bias()
+        ov_circuit = torch.einsum("onh,nhi->oni", reshaped_o, v_weight)
+        ov_bias_circuit = torch.einsum("onh,bnxh->bnxo", reshaped_o, v_bias)
+        return ov_circuit, ov_bias_circuit
+
+    def forward_ov_circuit(self, x, attn_mask=None):
+        B, N, C = x.shape
+        q_weight, k_weight, _ = self._split_qkv_weight()
+        q_bias, k_bias, _ = self._split_qkv_bias()
+        q = torch.einsum("bnc,hdc->bhnd", x, q_weight) + q_bias
+        k = torch.einsum("bnc,hdc->bhnd", x, k_weight) + k_bias
+        ov, ov_bias = self._get_ov_circuit()
+        v = torch.einsum("bnc,dhc->bhnd", x, ov) + ov_bias
+
+        dk = q.size()[-1]
+        q = q / math.sqrt(dk)
+        attn = q @ k.transpose(-2, -1)
+        if attn_mask is not None:
+            attn += attn_mask
+        attn = attn.softmax(dim=-1)
+        x = torch.einsum("bhnm,bhmc->bnmhc", attn, v)
+        x = x.sum(axis=[2, 3]) + self.out_proj.bias
+        return x
+
+    def forward(self, x, attn_mask=None, method: Text = "ov_circuit"):
+        if method == "direct":
+            return self.forward_direct(x, attn_mask=attn_mask)
+        elif method == "qkv":
+            return self.forward_qkv(x, attn_mask=attn_mask)
+        elif method == "head":
+            return self.forward_per_head(x, attn_mask=attn_mask)
+        elif method == "ov_circuit":
+            return self.forward_ov_circuit(x, attn_mask=attn_mask)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        d_model: int,
+        n_head: int,
+        mlp_ratio: float = 4.0,
+        ls_init_value: float = None,
+        act_layer: Callable = nn.GELU,
+        norm_layer: Callable = LayerNorm,
+        layer_id: Optional[int] = None,
+    ):
+        super().__init__()
+        self.ln_1 = norm_layer(d_model)
+        self.attn = MultiheadAttention(d_model, n_head)
+        self.layer_id = layer_id
+
+        self.ls_1 = (
+            LayerScale(d_model, ls_init_value)
+            if ls_init_value is not None
+            else nn.Identity()
+        )
+
+        self.ln_2 = norm_layer(d_model)
+        self.mlp_width = int(d_model * mlp_ratio)
+        self.mlp = MLP(
+            d_model, 
+            self.mlp_width, 
+            act_layer=act_layer,
+            layer_id=layer_id,
+        )
+        self.ls_2 = (
+            LayerScale(d_model, ls_init_value)
+            if ls_init_value is not None
+            else nn.Identity()
+        )
+
+    def attention(
+        self,
+        q_x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        method: Text = "direct",
+    ):
+        attn_mask = attn_mask.to(q_x.dtype) if attn_mask is not None else None
+        return self.attn(q_x, attn_mask=attn_mask, method=method)
+
+    def forward(
+        self,
+        q_x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        attn_method: Text = "direct",
+        neuron_dict=None,
+        num_register_tokens=0
+    ):
+        after_ln1 = self.ln_1(q_x)
+        after_attn = self.attention(
+            q_x=after_ln1, attn_mask=attn_mask, method=attn_method
+        )
+        x = q_x + self.ls_1(after_attn)
+        after_ln2 = self.ln_2(x)
+        after_mlp = self.mlp(after_ln2, neuron_dict=neuron_dict, num_register_tokens=num_register_tokens)
+        x = x + self.ls_2(after_mlp)
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(
+        self,
+        width: int,
+        layers: int,
+        heads: int,
+        mlp_ratio: float = 4.0,
+        ls_init_value: float = None,
+        act_layer: Callable = nn.GELU,
+        norm_layer: Callable = LayerNorm,
+    ):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.grad_checkpointing = False
+
+        self.resblocks = nn.ModuleList(
+            [
+                ResidualAttentionBlock(
+                    width,
+                    heads,
+                    mlp_ratio,
+                    ls_init_value=ls_init_value,
+                    act_layer=act_layer,
+                    norm_layer=norm_layer,
+                    layer_id=i,
+                )
+                for i in range(layers)
+            ]
+        )
+
+    def get_cast_dtype(self) -> torch.dtype:
+        if hasattr(self.resblocks[0].mlp.c_fc, "int8_original_dtype"):
+            return self.resblocks[0].mlp.c_fc.int8_original_dtype
+        return self.resblocks[0].mlp.c_fc.weight.dtype
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        attn_method: Text = "direct",
+        neuron_dict=None,
+        num_register_tokens=0
+    ):
+        for r in self.resblocks:
+            if self.grad_checkpointing and not torch.jit.is_scripting():
+                raise ValueError("grad_checkpointing not implemented")
+            else:
+                x = r(
+                    x, 
+                    attn_mask=attn_mask, 
+                    attn_method=attn_method,
+                    neuron_dict=neuron_dict,
+                    num_register_tokens=num_register_tokens
+                )
+        return x
+
+
+class VisionTransformer(nn.Module):
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+        self,
+        image_size: int,
+        patch_size: int,
+        width: int,
+        layers: int,
+        heads: int,
+        mlp_ratio: float,
+        ls_init_value: float = None,
+        global_average_pool: bool = False,
+        attentional_pool: bool = False,
+        n_queries: int = 256,
+        attn_pooler_heads: int = 8,
+        output_dim: int = 512,
+        patch_dropout: float = 0.0,
+        input_patchnorm: bool = False,
+        act_layer: Callable = nn.GELU,
+        norm_layer: Callable = LayerNorm,
+        output_tokens: bool = False,
+    ):
+        super().__init__()
+        self.output_tokens = output_tokens
+        image_height, image_width = self.image_size = to_2tuple(image_size)
+        patch_height, patch_width = self.patch_size = to_2tuple(patch_size)
+        self.grid_size = (image_height // patch_height, image_width // patch_width)
+        self.output_dim = output_dim
+
+        self.num_register_tokens = 0
+        self.neuron_dict = None
+
+        self.input_patchnorm = input_patchnorm
+
+        if input_patchnorm:
+            patch_input_dim = patch_height * patch_width * 3
+            self.patchnorm_pre_ln = LayerNorm(patch_input_dim)
+            self.conv1 = nn.Linear(patch_input_dim, width)
+        else:
+            self.patchnorm_pre_ln = nn.Identity()
+            self.conv1 = nn.Conv2d(
+                in_channels=3,
+                out_channels=width,
+                kernel_size=patch_size,
+                stride=patch_size,
+                bias=False,
+            )
+
+        scale = width**-0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(
+            scale * torch.randn(self.grid_size[0] * self.grid_size[1] + 1, width)
+        )
+
+        self.width = width
+        self.scale = scale
+        self.extra_token = self.scale * torch.randn(width)
+
+        self.patch_dropout = (
+            PatchDropout(patch_dropout) if patch_dropout > 0.0 else nn.Identity()
+        )
+
+        self.ln_pre = norm_layer(width)
+        self.transformer = Transformer(
+            width,
+            layers,
+            heads,
+            mlp_ratio,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+
+        self.global_average_pool = global_average_pool
+        if attentional_pool:
+            self.attn_pool = AttentionalPooler(
+                output_dim, width, n_head=attn_pooler_heads, n_queries=n_queries
+            )
+            self.ln_post = norm_layer(output_dim)
+            self.proj = nn.Parameter(scale * torch.randn(output_dim, output_dim))
+        else:
+            self.attn_pool = None
+            self.ln_post = norm_layer(width)
+            self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.grad_checkpointing = enable
+
+    def _global_pool(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.global_average_pool:
+            return x.mean(dim=1), x
+        else:
+            return x[:, 0], x[:, 1:]
+
+    def forward(self, x: torch.Tensor, attn_method: Text = "direct", num_register_tokens = None, neuron_dict=None):
+        # to patches
+
+        if num_register_tokens is None and neuron_dict is None:
+            num_register_tokens = self.num_register_tokens
+            neuron_dict = self.neuron_dict
+
+        if self.input_patchnorm:
+            x = x.reshape(
+                x.shape[0],
+                x.shape[1],
+                self.grid_size[0],
+                self.patch_size[0],
+                self.grid_size[1],
+                self.patch_size[1],
+            )
+            x = x.permute(0, 2, 4, 1, 3, 5)
+            x = x.reshape(x.shape[0], self.grid_size[0] * self.grid_size[1], -1)
+            x = self.patchnorm_pre_ln(x)
+            x = self.conv1(x)
+        else:
+            x = self.conv1(x)
+            x = x.reshape(x.shape[0], x.shape[1], -1)
+            x = x.permute(0, 2, 1)
+
+        # class embeddings and positional embeddings
+        x = torch.cat([
+                self.class_embedding.to(x.dtype)
+                + torch.zeros(
+                    x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device
+                ),
+                x,
+            ],
+            dim=1,
+        )
+        x = x + self.positional_embedding.to(x.dtype)
+
+        extra_token_embeddings = []
+        total_patches = x.shape[1] - 1
+        for i in range(num_register_tokens):
+            extra_token_embeddings.append(
+                torch.zeros(
+                    x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device
+                ),
+            )
+            
+        # Add extra tokens
+        if num_register_tokens > 0:
+            x = torch.cat([x, *extra_token_embeddings], dim=1)
+
+        x = self.patch_dropout(x)
+        x = self.ln_pre(x)
+        
+        x = self.transformer(x, attn_mask=None, attn_method=attn_method, neuron_dict=neuron_dict, num_register_tokens=num_register_tokens)
+        
+        if self.attn_pool is not None:
+            x = self.attn_pool(x)
+            x = self.ln_post(x)
+            pooled, tokens = self._global_pool(x)
+        else:
+            pooled, tokens = self._global_pool(x)
+            pooled = self.ln_post(pooled)
+
+        if self.proj is not None:
+            pooled = pooled @ self.proj
+
+        if self.output_tokens:
+            return pooled, tokens
+
+        return pooled
+
+
+class TextTransformer(nn.Module):
+    output_tokens: torch.jit.Final[bool]
+
+    def __init__(
+        self,
+        context_length: int = 77,
+        vocab_size: int = 49408,
+        width: int = 512,
+        heads: int = 8,
+        layers: int = 12,
+        ls_init_value: float = None,
+        output_dim: int = 512,
+        act_layer: Callable = nn.GELU,
+        norm_layer: Callable = LayerNorm,
+        embed_cls: bool = False,
+        pad_id: int = 0,
+        output_tokens: bool = False,
+    ):
+        super().__init__()
+        self.output_tokens = output_tokens
+        self.num_pos = self.context_length = context_length
+        self.vocab_size = vocab_size
+        self.width = width
+        self.output_dim = output_dim
+        self.heads = heads
+        self.pad_id = pad_id
+
+        self.text_projection = nn.Parameter(torch.empty(width, output_dim))
+
+        if embed_cls:
+            self.cls_emb = nn.Parameter(torch.empty(width))
+            self.num_pos += 1
+        else:
+            self.cls_emb = None
+
+        self.token_embedding = nn.Embedding(vocab_size, width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.num_pos, width))
+        self.transformer = Transformer(
+            width=width,
+            layers=layers,
+            heads=heads,
+            ls_init_value=ls_init_value,
+            act_layer=act_layer,
+            norm_layer=norm_layer,
+        )
+        self.ln_final = norm_layer(width)
+
+        self.register_buffer("attn_mask", self.build_attention_mask(), persistent=False)
+
+        self.init_parameters()
+
+    def init_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+        if self.cls_emb is not None:
+            nn.init.normal_(self.cls_emb, std=0.01)
+
+        proj_std = (self.transformer.width**-0.5) * (
+            (2 * self.transformer.layers) ** -0.5
+        )
+        attn_std = self.transformer.width**-0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width**-0.5)
+
+    @torch.jit.ignore
+    def set_grad_checkpointing(self, enable=True):
+        self.transformer.grad_checkpointing = enable
+
+    def build_attention_mask(self):
+        mask = torch.empty(self.num_pos, self.num_pos)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)
+        return mask
+
+    def build_cls_mask(self, text, cast_dtype: torch.dtype):
+        cls_mask = (text != self.pad_id).unsqueeze(1)
+        cls_mask = F.pad(cls_mask, (1, 0, cls_mask.shape[2], 0), value=1.0)
+        additive_mask = torch.empty(
+            cls_mask.shape, dtype=cast_dtype, device=cls_mask.device
+        )
+        additive_mask.fill_(0)
+        additive_mask.masked_fill_(~cls_mask, float("-inf"))
+        additive_mask = torch.repeat_interleave(additive_mask, self.heads, 0)
+        return additive_mask
+
+    def _repeat(self, t, N: int):
+        return t.reshape(1, 1, -1).repeat(N, 1, 1)
+
+    def forward(self, text, attn_method: Text = "direct"):
+        cast_dtype = self.transformer.get_cast_dtype()
+        seq_len = text.shape[1]
+
+        x = self.token_embedding(text).to(cast_dtype)
+        attn_mask = self.attn_mask
+        if self.cls_emb is not None:
+            seq_len += 1
+            x = torch.cat([x, self._repeat(self.cls_emb, x.shape[0])], dim=1)
+            cls_mask = self.build_cls_mask(text, cast_dtype)
+            attn_mask = (
+                attn_mask[None, :seq_len, :seq_len] + cls_mask[:, :seq_len, :seq_len]
+            )
+
+        x = x + self.positional_embedding[:seq_len].to(cast_dtype)
+        x = self.transformer(x, attn_mask=attn_mask, attn_method=attn_method)
+
+        if self.cls_emb is not None:
+            pooled, tokens = x[:, -1], x[:, :-1]
+            pooled = self.ln_final(pooled)
+        else:
+            x = self.ln_final(x)
+            pooled, tokens = x[torch.arange(x.shape[0]), text.argmax(dim=-1)], x
+
+        if self.text_projection is not None:
+            pooled = pooled @ self.text_projection
+
+        if self.output_tokens:
+            return pooled, tokens
+
+        return pooled

utils.py

@@ -0,0 +1,34 @@
+import time
+import torch
+from datetime import datetime
+import os
+
+class SaveOcassionally:
+  def __init__(self, out, every_sec = None, every_count = None):
+    assert every_sec != None or every_count != None
+
+    self.out = out
+    self.curr_time = time.time()
+    self.every_sec = every_sec
+    self.cnt = 0
+    self.every_count = every_count
+    self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    if "TIMESTAMP" in self.out:
+      self.out = self.out.replace("TIMESTAMP", self.timestamp)
+      print(f"Replacing TIMESTAMP with {self.timestamp}")
+
+    # Ensure the directory exists
+    out_dir = os.path.abspath(os.path.dirname(self.out))
+    if not os.path.exists(out_dir):
+      os.makedirs(out_dir)
+
+  def save(self, obj):
+    self.cnt += 1
+    if self.every_sec != None and time.time() - self.curr_time > self.every_sec:
+      torch.save(obj, self.out)
+    elif self.every_count != None and self.cnt % self.every_count == 0:
+      torch.save(obj, self.out)
+
+  def force_save(self, obj):
+    torch.save(obj, self.out)

utils/utils.py

@@ -0,0 +1,34 @@
+import time
+import torch
+from datetime import datetime
+import os
+
+class SaveOcassionally:
+  def __init__(self, out, every_sec = None, every_count = None):
+    assert every_sec != None or every_count != None
+
+    self.out = out
+    self.curr_time = time.time()
+    self.every_sec = every_sec
+    self.cnt = 0
+    self.every_count = every_count
+    self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+    if "TIMESTAMP" in self.out:
+      self.out = self.out.replace("TIMESTAMP", self.timestamp)
+      print(f"Replacing TIMESTAMP with {self.timestamp}")
+
+    # Ensure the directory exists
+    out_dir = os.path.abspath(os.path.dirname(self.out))
+    if not os.path.exists(out_dir):
+      os.makedirs(out_dir)
+
+  def save(self, obj):
+    self.cnt += 1
+    if self.every_sec != None and time.time() - self.curr_time > self.every_sec:
+      torch.save(obj, self.out)
+    elif self.every_count != None and self.cnt % self.every_count == 0:
+      torch.save(obj, self.out)
+
+  def force_save(self, obj):
+    torch.save(obj, self.out)

vocab/bpe_simple_vocab_16e6.txt.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a
+size 1356917

zeroshot_classifier.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7dff47ac37ed4b67771bf6cf651a55dcf95d22eddc91acce2f54638ec82c6783
+size 1537240