File size: 5,323 Bytes
1bd7d53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | ---
license: apache-2.0
language:
- en
tags:
- vision-language
- image-text
- vit
- gpt2
- contrastive-learning
- siglip
- sigmoid-loss
---
# SigLIP — vit_base_patch16_224 / DataComp-large
Checkpoint for **SigLIP**, a contrastive vision-language pretraining baseline using the pairwise sigmoid (SigLIP) loss over a BN-MLP projection head.
Trained on DataComp-large for 200,000 steps with batch size 4,096.
## Model summary
| Property | Value |
|---|---|
| Method | SigLIP |
| Vision encoder | `vit_base_patch16_224` (timm) |
| Text encoder | GPT-2 (12L / 12H / 768D) |
| Embedding dim | 512 |
| Projection head | Linear→BN→GELU→Linear (width 2048) |
| Training objective | Pairwise sigmoid (SigLIP) contrastive loss |
| Training data | DataComp-large |
| Training steps | 200,000 |
## Usage
```python
import torch
import timm
from transformers import GPT2Config, GPT2Model, AutoTokenizer
from safetensors.torch import load_file
from torchvision.ops import MLP
import torch.nn as nn
HIDDEN = 768
EMBED = 512
# ── Vision encoder ───────────────────────────────────────────────────────────
vision_encoder = timm.create_model(
"vit_base_patch16_224", pretrained=False, num_classes=0, dynamic_img_size=True
)
vision_pre_proj = nn.Sequential(
nn.Linear(HIDDEN, 2048), nn.BatchNorm1d(2048), nn.GELU(), nn.Linear(2048, EMBED)
)
# ── Text encoder ─────────────────────────────────────────────────────────────
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token
def tokenize_with_eos_readout(tokenizer, text, max_length=77):
ids = tokenizer(
text,
add_special_tokens=False,
truncation=True,
max_length=max_length - 1,
)["input_ids"] + [tokenizer.eos_token_id]
pad_len = max_length - len(ids)
input_ids = torch.tensor([ids + [tokenizer.pad_token_id] * pad_len])
attention_mask = torch.tensor([[1] * len(ids) + [0] * pad_len])
return dict(input_ids=input_ids, attention_mask=attention_mask)
def last_unmasked_token(hidden, attention_mask):
lengths = attention_mask.sum(dim=1).clamp(min=1).long()
gather_idx = (lengths - 1).view(-1, 1, 1).expand(-1, 1, hidden.size(-1))
return hidden.gather(1, gather_idx).squeeze(1)
text_encoder = GPT2Model(GPT2Config(
n_embd=HIDDEN, n_layer=12, n_head=12,
n_inner=HIDDEN * 4, vocab_size=tokenizer.vocab_size,
attn_pdrop=0.0, resid_pdrop=0.0, embd_pdrop=0.0,
))
text_pre_proj = nn.Sequential(
nn.Linear(HIDDEN, 2048), nn.BatchNorm1d(2048), nn.GELU(), nn.Linear(2048, EMBED)
)
# ── Load weights ─────────────────────────────────────────────────────────────
from huggingface_hub import hf_hub_download
vision_weights = load_file(hf_hub_download("lukaskuhndkfz/SigLIP-ViT-B-DataComp-200k", "vision_encoder.safetensors"))
text_weights = load_file(hf_hub_download("lukaskuhndkfz/SigLIP-ViT-B-DataComp-200k", "text_encoder.safetensors"))
vision_encoder.load_state_dict({k[len("encoder."):]: v for k, v in vision_weights.items() if k.startswith("encoder.")})
vision_pre_proj.load_state_dict({k[len("pre_proj."):]: v for k, v in vision_weights.items() if k.startswith("pre_proj.")})
text_encoder.load_state_dict({k[len("encoder."):]: v for k, v in text_weights.items() if k.startswith("encoder.")})
text_pre_proj.load_state_dict({k[len("pre_proj."):]: v for k, v in text_weights.items() if k.startswith("pre_proj.")})
vision_encoder.eval()
text_encoder.eval()
# ── Encode an image ──────────────────────────────────────────────────────────
from torchvision import transforms
from PIL import Image
transform = transforms.Compose([
transforms.Resize(224), transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
image = Image.open("image.jpg").convert("RGB")
pixel_values = transform(image).unsqueeze(0)
with torch.no_grad():
image_features = vision_pre_proj(vision_encoder(pixel_values)) # (1, 512)
# ── Encode a caption ─────────────────────────────────────────────────────────
inputs = tokenize_with_eos_readout(tokenizer, "a photo of a cat")
with torch.no_grad():
hidden = text_encoder(**inputs).last_hidden_state
text_hidden = last_unmasked_token(hidden, inputs["attention_mask"])
text_features = text_pre_proj(text_hidden) # (1, 512)
```
## Files
| File | Contents |
|---|---|
| `vision_encoder.safetensors` | Vision encoder (`encoder.*`), pre-projection head (`pre_proj.*`), and cross-modal projector MLP (`projector.*`) |
| `text_encoder.safetensors` | Text encoder (`encoder.*`), pre-projection head (`pre_proj.*`), and cross-modal projector MLP (`projector.*`) |
| `config.json` | Architecture and training hyperparameters |
|