Add Geometric Parameterization (GmP) for CLIP ft

Extremely experimental Geometric Parameterization (GmP): .weight -> .theta, .r

Author: zer0int

exp-ft-B-GmP-finetune-OpenAI-ViT-L-14.py

@@ -0,0 +1,63 @@
+import torch
+import torch.nn.functional as F
+from gmpclip.modeloriginal import CLIP  # Make sure this imports the original CLIP class definition
+
+def convert_back_to_original(state_dict):
+    new_state_dict = {}
+    
+    for key, value in state_dict.items():
+        if key.endswith(".theta"):
+            base_key = key.replace(".theta", "")
+            r_key = base_key + ".r"
+            new_weight = state_dict[r_key] * F.normalize(value, p=2, dim=1)
+            new_state_dict[base_key + ".weight"] = new_weight
+        elif key.endswith(".r") or key.endswith(".theta"):
+            continue  # Skip the .r and .theta keys
+        else:
+            new_state_dict[key] = value
+
+    return new_state_dict
+
+# Example usage
+# Load the fine-tuned model object
+modelft = torch.load("ft-checkpoints/clip_ft_20.pt")
+
+# Extract model parameters from the fine-tuned model
+embed_dim = modelft.text_projection.shape[1]
+image_resolution = modelft.visual.input_resolution
+vision_layers = modelft.visual.transformer.layers
+vision_width = modelft.visual.conv1.out_channels
+vision_patch_size = modelft.visual.conv1.kernel_size[0]
+context_length = modelft.context_length
+vocab_size = modelft.vocab_size
+transformer_width = modelft.transformer.width
+transformer_heads = modelft.transformer.resblocks[0].attn.num_heads
+transformer_layers = modelft.transformer.layers
+
+# Convert the fine-tuned model to a state_dict
+fine_tuned_state_dict = modelft.state_dict()
+
+# Convert back to original weights
+original_state_dict = convert_back_to_original(fine_tuned_state_dict)
+
+# Rebuild the original model using the converted state_dict
+original_model = CLIP(
+    embed_dim=embed_dim,
+    image_resolution=image_resolution,
+    vision_layers=vision_layers,
+    vision_width=vision_width,
+    vision_patch_size=vision_patch_size,
+    context_length=context_length,
+    vocab_size=vocab_size,
+    transformer_width=transformer_width,
+    transformer_heads=transformer_heads,
+    transformer_layers=transformer_layers
+)
+
+# Load the converted state_dict into the original model
+original_model.load_state_dict(original_state_dict)
+
+# Save the original model object
+torch.save(original_model, "ft-checkpoints/full_model_converted_model.pth")
+
+print("Model has been successfully converted back to the original format and saved.")

exp-ft-C-convert-GmP-back-to-weight.py

@@ -0,0 +1,75 @@
+import torch
+import gmpclip as clip
+import matplotlib.pyplot as plt
+import seaborn as sns
+import concurrent.futures
+from itertools import chain
+import gc
+import numpy as np
+
+def load_model_checkpoint(checkpoint_path):
+    model = torch.load(checkpoint_path)
+    return model
+
+def extract_geometric_params(model):
+    geometric_params = {'theta': [], 'r': []}
+    for name, param in model.named_parameters():
+        if name.endswith('.theta'):
+            geometric_params['theta'].append(param.detach().cpu().numpy())
+        elif name.endswith('.r'):
+            geometric_params['r'].append(param.detach().cpu().numpy())
+    return geometric_params
+
+# Downsampling factor - without, this takes >>100 GB RAM + an hour or so. Factor 10 => 5 minutes (Ryzen 9).
+# Adjust factor as needed.
+def downsample(data, factor=10):
+    return data[::factor]
+
+def visualize_params(params, title, iteration):
+    theta_flat = list(chain.from_iterable([item.flatten() for sublist in params['theta'] for item in sublist]))
+    r_flat = list(chain.from_iterable([item.flatten() for sublist in params['r'] for item in sublist]))
+    
+    # Downsample data to reduce memory usage
+    theta_flat = downsample(theta_flat)
+    r_flat = downsample(r_flat)
+    
+    plt.figure(figsize=(12, 6))
+    
+    plt.subplot(1, 2, 1)
+    sns.histplot(theta_flat, bins=50, kde=True)
+    plt.title('Distribution of Theta Components')
+    
+    plt.subplot(1, 2, 2)
+    sns.histplot(r_flat, bins=50, kde=True)
+    plt.title('Distribution of R Components')
+    
+    plt.suptitle(title)
+    plt.savefig(f'geometric_params_visualization_{iteration}.png')
+    plt.close()
+
+def process_checkpoint(checkpoint, iteration):
+    model = load_model_checkpoint(checkpoint)
+    params = extract_geometric_params(model)
+    visualize_params(params, f'Checkpoint: {checkpoint}', iteration)
+    # Explicitly call garbage collection
+    gc.collect()
+
+if __name__ == '__main__':
+    # List of checkpoints - fine-tuned model saves:
+    checkpoints = ["ft-checkpoints/clip_ft_5.pt", "ft-checkpoints/clip_ft_10.pt", 
+                   "ft-checkpoints/clip_ft_15.pt", "ft-checkpoints/clip_ft_20.pt"]
+
+    # Split the list of checkpoints into smaller batches if 4 at once consumes too much RAM (RAM, not VRAM!)
+    batch_size = 4
+    batches = [checkpoints[i:i + batch_size] for i in range(0, len(checkpoints), batch_size)]
+
+    for batch in batches:
+        # Parallel processing using concurrent.futures
+        with concurrent.futures.ProcessPoolExecutor() as executor:
+            futures = [executor.submit(process_checkpoint, checkpoint, i+1) for i, checkpoint in enumerate(batch)]
+        # Wait for all futures to complete
+        concurrent.futures.wait(futures)
+        # Explicitly call garbage collection
+        gc.collect()
+
+    print("All visualizations have been generated and saved.")

exp-ft-X-visualize-theta-r-barplots.py

@@ -0,0 +1 @@
+from .clip import *

gmpclip/__init__.py

@@ -0,0 +1,247 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Any, Union, List
+from pkg_resources import packaging
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+from tqdm import tqdm
+
+from .model import build_model
+from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+if packaging.version.parse(torch.__version__) < packaging.version.parse("1.7.1"):
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", "tokenize"]
+_tokenizer = _Tokenizer()
+
+_MODELS = {
+    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
+    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+def _download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    # Commenting out the expected_sha256 as we're bypassing checksum verification
+    # expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, 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):
+        # Bypassing the SHA256 checksum verification
+        return download_target
+             
+                                                                                                                       
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+                                                                                         
+                                                                                                  
+
+    # Bypassing the SHA256 checksum verification on download completion
+    return download_target
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+    """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
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    jit : bool
+        Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/clip"
+
+    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 name in _MODELS:
+        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    with open(model_path, 'rb') as opened_file:
+        try:
+            # loading JIT archive
+            #model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
+            model = torch.load(opened_file, map_location=device if jit else "cpu").eval()
+            state_dict = None
+        except RuntimeError:
+            # loading saved state dict
+            if jit:
+                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+                jit = False
+            state_dict = torch.load(opened_file, map_location="cpu")
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).to(device)
+        if str(device) == "cpu":
+            model.float()
+        return model, _transform(model.visual.input_resolution)
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def _node_get(node: torch._C.Node, key: str):
+        """Gets attributes of a node which is polymorphic over return type.
+        
+        From https://github.com/pytorch/pytorch/pull/82628
+        """
+        sel = node.kindOf(key)
+        return getattr(node, sel)(key)
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(_node_get(node, "value")).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if _node_get(inputs[i].node(), "value") == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model, _transform(model.input_resolution.item())
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, 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
+
+    truncate: bool
+        Whether to truncate the text in case its encoding is longer than the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
+    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<|startoftext|>"]
+    eot_token = _tokenizer.encoder["<|endoftext|>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"):
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+    else:
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            if truncate:
+                tokens = tokens[:context_length]
+                tokens[-1] = eot_token
+            else:
+                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result