[Feature] Add SigLIP loss and forward template

benchmarks/diffusiondb2m/siglip-vitb.py

@@ -0,0 +1,220 @@
+import argparse
+from functools import partial
+
+import lightning as pl
+import torch
+import torch.nn.functional as F
+from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint
+from lightning.pytorch.loggers import WandbLogger
+from transformers import (
+    AutoImageProcessor,
+    AutoTokenizer,
+    SiglipTextModel,
+    SiglipVisionModel,
+)
+
+import stable_pretraining as spt
+from stable_pretraining import forward
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--lr", type=float, default=0.001)
+parser.add_argument("--num_devices", type=int, default=8)
+parser.add_argument("--global_batch", type=int, default=4096)
+parser.add_argument("--num_epochs", type=int, default=8)
+parser.add_argument("--val_percent", type=float, default=0.10)
+parser.add_argument("--resume_ckpt_path", type=str, default=None)
+args = parser.parse_args()
+
+lr = args.lr
+num_devices = args.num_devices
+global_batch = args.global_batch
+batch_size = global_batch // num_devices
+num_epochs = args.num_epochs
+val_percent = args.val_percent
+resume_ckpt_path = args.resume_ckpt_path
+
+model_name = "google/siglip-base-patch16-224"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+image_processor = AutoImageProcessor.from_pretrained(model_name)
+vision_model = SiglipVisionModel.from_pretrained(model_name)
+text_model = SiglipTextModel.from_pretrained(model_name)
+
+
+def tokenize(text: str, tokenizer: AutoTokenizer):
+    data = tokenizer(
+        text,
+        return_tensors="pt",
+        padding="max_length",
+        max_length=tokenizer.model_max_length,
+        truncation=True,
+    )
+    return data["input_ids"].squeeze(0), data["attention_mask"].squeeze(0)
+
+
+image_transform = spt.data.transforms.Compose(
+    spt.data.transforms.Resize((224, 224)),
+    spt.data.transforms.ToImage(
+        mean=image_processor.image_mean,
+        std=image_processor.image_std,
+    ),
+    spt.data.transforms.LambdaTransform(
+        fn=partial(tokenize, tokenizer=tokenizer),
+        source="prompt",
+        targets=("tokenized_prompt", "attention_mask"),
+    ),
+)
+
+
+train_base = spt.data.HFDataset(
+    "poloclub/diffusiondb",
+    "2m_all",
+    split="train",
+    transform=image_transform,
+    remove_columns=[
+        "timestamp",
+        "user_name",
+        "prompt_nsfw",
+        "image_nsfw",
+        "sampler",
+    ],
+)
+
+size = len(train_base)
+val_n = int(size * val_percent)
+val_dataset = spt.data.Subset(train_base, range(0, val_n))
+train_dataset = spt.data.Subset(train_base, range(val_n, size))
+
+
+train_dataloader = torch.utils.data.DataLoader(
+    dataset=train_dataset,
+    batch_size=batch_size,
+    num_workers=16,
+    shuffle=True,
+    drop_last=True,
+    pin_memory=True,
+    persistent_workers=True,
+    prefetch_factor=4,
+)
+val_dataloader = torch.utils.data.DataLoader(
+    dataset=val_dataset,
+    batch_size=batch_size,
+    num_workers=8,
+    shuffle=False,
+    pin_memory=True,
+    persistent_workers=True,
+    prefetch_factor=4,
+)
+
+data = spt.data.DataModule(train=train_dataloader, val=val_dataloader)
+
+
+class SigLIPMonitor(pl.Callback):
+    """Log retrieval and pairwise sigmoid statistics for SigLIP training."""
+
+    def __init__(self, log_every_n_steps: int = 10):
+        super().__init__()
+        self.every = log_every_n_steps
+
+    @torch.no_grad()
+    def _log(self, trainer: pl.Trainer, pl_module, outputs: dict, stage: str):
+        img = F.normalize(outputs["image_embeds"], dim=-1)
+        txt = F.normalize(outputs["text_embeds"], dim=-1)
+
+        loss_fn = pl_module.siglip_loss
+        logits = loss_fn.logit_scale.exp() * (img @ txt.T) + loss_fn.logit_bias
+        batch_size = logits.size(0)
+        diag = torch.arange(batch_size, device=logits.device)
+
+        r1_i2t = (logits.argmax(dim=1) == diag).float().mean()
+        r1_t2i = (logits.argmax(dim=0) == diag).float().mean()
+        pos_prob = torch.sigmoid(logits[diag, diag]).mean()
+        cos_pos = F.cosine_similarity(img, txt, dim=-1).mean()
+
+        metrics = {
+            f"{stage}/retrieval/R@1_i2t": float(r1_i2t.cpu()),
+            f"{stage}/retrieval/R@1_t2i": float(r1_t2i.cpu()),
+            f"{stage}/contrast/pos_prob": float(pos_prob.cpu()),
+            f"{stage}/align/cos_pos": float(cos_pos.cpu()),
+            f"{stage}/config/logit_scale": float(loss_fn.logit_scale.exp().cpu()),
+            f"{stage}/config/logit_bias": float(loss_fn.logit_bias.cpu()),
+        }
+        if batch_size > 1:
+            neg = logits.masked_fill(
+                torch.eye(batch_size, dtype=torch.bool, device=logits.device),
+                float("-inf"),
+            )
+            top_neg = neg.max(dim=1).values
+            margin = logits[diag, diag] - top_neg
+            metrics[f"{stage}/contrast/margin"] = float(margin.mean().cpu())
+
+        trainer.logger.log_metrics(metrics, step=trainer.global_step)
+
+    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
+        if trainer.global_step % self.every == 0:
+            self._log(trainer, pl_module, outputs, "train")
+
+    def on_validation_batch_end(
+        self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0
+    ):
+        self._log(trainer, pl_module, outputs, "val")
+
+
+module = spt.Module(
+    vision_model=vision_model,
+    text_model=text_model,
+    forward=forward.siglip_forward,
+    siglip_loss=spt.losses.SigLIPLoss(),
+    optim={
+        "optimizer": {
+            "type": "AdamW",
+            "lr": lr,
+            "weight_decay": 1.0e-6,
+            "betas": (0.9, 0.98),
+        },
+        "scheduler": {
+            "type": "LinearWarmupCosineAnnealing",
+            "total_steps": (len(train_dataloader) // num_devices) * num_epochs,
+            "peak_step": 0.1,
+        },
+        "interval": "step",
+    },
+)
+
+wandb_logger = WandbLogger(
+    entity="stable-pretraining",
+    project="diffusiondb2m-siglip",
+    name="siglip-vit-b16-diffusiondb2m-32k",
+    log_model=False,
+)
+
+trainer = pl.Trainer(
+    max_epochs=num_epochs,
+    num_sanity_val_steps=0,
+    callbacks=[
+        ModelCheckpoint(
+            monitor="fit/loss_step",
+            mode="min",
+            every_n_epochs=1,
+            save_top_k=-1,
+            dirpath="/your/path/to/checkpoints",
+        ),
+        LearningRateMonitor(logging_interval="step"),
+        SigLIPMonitor(log_every_n_steps=10),
+    ],
+    precision="bf16-mixed",
+    logger=wandb_logger,
+    enable_checkpointing=True,
+    devices=num_devices,
+    accelerator="gpu",
+    strategy="ddp",
+)
+
+manager = spt.Manager(
+    trainer=trainer,
+    module=module,
+    data=data,
+    ckpt_path=resume_ckpt_path,
+)
+
+manager()

docs/source/api/forward.rst

@@ -247,6 +247,34 @@ Barlow Twins
         _target_: stable_pretraining.losses.BarlowTwinsLoss
         lambda_: 0.005
 
+SigLIP
+~~~~~~
+
+.. autofunction:: siglip_forward
+
+**Required Module Attributes:**
+
+- ``vision_model``: Image encoder returning ``image_embeds`` or ``pooler_output``
+- ``text_model``: Text encoder returning ``text_embeds`` or ``pooler_output``
+- ``siglip_loss``: SigLIP loss function
+
+**Expected Batch Keys:**
+
+- ``image``: Image tensor passed to the vision model
+- ``tokenized_prompt``: Token ids passed to the text model
+- ``attention_mask``: Optional text attention mask
+
+**Example Config:**
+
+.. code-block:: yaml
+
+    module:
+      forward: stable_pretraining.forward.siglip_forward
+      vision_model: ...
+      text_model: ...
+      siglip_loss:
+        _target_: stable_pretraining.losses.SigLIPLoss
+
 Supervised
 ~~~~~~~~~~
 

docs/source/api/losses.rst

@@ -12,3 +12,5 @@ stable_pretraining.losses
    NegativeCosineSimilarity
    VICRegLoss
    BarlowTwinsLoss
+   CLIPLoss
+   SigLIPLoss

docs/source/references.bib

@@ -96,6 +96,13 @@ @inproceedings{radford2021learning
   organization={PmLR}
 }
 
+@article{zhai2023sigmoid,
+  title={Sigmoid loss for language image pre-training},
+  author={Zhai, Xiaohua and Mustafa, Basil and Kolesnikov, Alexander and Beyer, Lucas},
+  journal={arXiv preprint arXiv:2303.15343},
+  year={2023}
+}
+
 @article{caron2020unsupervised,
   title={Unsupervised learning of visual features by contrasting cluster assignments},
   author={Caron, Mathilde and Misra, Ishan and Mairal, Julien and Goyal, Priya and Bojanowski, Piotr and Joulin, Armand},

stable_pretraining/forward.py

@@ -1074,3 +1074,66 @@ def dinov2_forward(self, batch, stage):
     )
 
     return out
+
+
+def _get_embedding_output(outputs, primary_key: str, fallback_key: str):
+    embedding = getattr(outputs, primary_key, None)
+    if embedding is None:
+        embedding = getattr(outputs, fallback_key, None)
+    if embedding is None:
+        raise ValueError(
+            f"Expected model output to expose '{primary_key}' or '{fallback_key}'."
+        )
+    return embedding
+
+
+def siglip_forward(self, batch, stage):
+    """Forward function for SigLIP image-text pretraining.
+
+    SigLIP learns aligned image-text representations with an independent
+    sigmoid loss over every image-text pair in the batch. Matching pairs are
+    expected to be aligned along the batch diagonal.
+
+    Args:
+        self: Module instance with ``vision_model``, ``text_model``, and
+            ``siglip_loss`` attributes. ``vision_model`` must return
+            ``image_embeds`` or ``pooler_output``; ``text_model`` must return
+            ``text_embeds`` or ``pooler_output``.
+        batch: Paired image-text batch dictionary. Must contain ``image`` and
+            ``tokenized_prompt``. May contain ``attention_mask``.
+        stage: Training stage ('train', 'val', or 'test')
+
+    Returns:
+        Dictionary containing normalized ``image_embeds`` and ``text_embeds``.
+        During training, also contains ``loss``.
+    """
+    out = {}
+
+    vision_outputs = self.vision_model(pixel_values=batch["image"])
+    image_embeds = _get_embedding_output(
+        vision_outputs, "image_embeds", "pooler_output"
+    )
+    image_embeds = torch.nn.functional.normalize(image_embeds, dim=-1)
+
+    text_outputs = self.text_model(
+        input_ids=batch["tokenized_prompt"],
+        attention_mask=batch.get("attention_mask"),
+    )
+    text_embeds = _get_embedding_output(text_outputs, "text_embeds", "pooler_output")
+    text_embeds = torch.nn.functional.normalize(text_embeds, dim=-1)
+
+    out["image_embeds"] = image_embeds
+    out["text_embeds"] = text_embeds
+
+    if self.training:
+        out["loss"] = self.siglip_loss(image_embeds, text_embeds)
+        self.log(
+            f"{stage}/loss",
+            out["loss"],
+            on_step=True,
+            on_epoch=True,
+            sync_dist=True,
+            prog_bar=True,
+        )
+
+    return out

stable_pretraining/losses/__init__.py

@@ -19,7 +19,7 @@
 )
 
 # Multimodal losses
-from .multimodal import CLIPLoss
+from .multimodal import CLIPLoss, SigLIPLoss
 
 # Reconstruction losses
 from .reconstruction import mae
@@ -43,6 +43,7 @@
     "BarlowTwinsLoss",
     "NTXEntLoss",
     "CLIPLoss",
+    "SigLIPLoss",
     # Reconstruction
     "mae",
     # Utils