Upload 12 files

src/__init__.py

@@ -0,0 +1 @@
+# Makes this directory a python package

src/evaluate.py

@@ -0,0 +1,192 @@
+# eval.py
+import os
+import json
+from tqdm import tqdm
+import torch
+from nltk.translate.bleu_score import corpus_bleu, SmoothingFunction
+from src.utils import count_encoder_decoder_params, load_experiment
+from src.inference import load_image, generate_caption
+from PIL import Image
+
+
+# COCO metrics
+try:
+    from pycocoevalcap.cider.cider import Cider
+    from pycocoevalcap.rouge.rouge import Rouge
+    HAS_COCOEVAL = True
+except ImportError:
+    print("WARNING: pycocoevalcap not installed → CIDEr/ROUGE disabled.")
+    HAS_COCOEVAL = False
+
+
+def evaluate(model, tokenizer, preprocess, image_size, data_dir="data/processed", save_dir="checkpoints", device="cuda"):
+
+
+    captions_path = os.path.join(data_dir, "captions.json")
+    splits_path = os.path.join(data_dir, "splits.json")
+
+    captions = json.load(open(captions_path))
+    splits = json.load(open(splits_path))
+    val_ids = splits["val"]
+
+    preds = []
+    refs_tokenized = []   # for BLEU
+    refs_strings = []     # for JSON log
+
+    print(f"Running evaluation on {len(val_ids)} images…\n")
+    with torch.no_grad():
+        for idx, img_id in enumerate(tqdm(val_ids, desc="Evaluating")):
+            img_path = os.path.join(data_dir, "images", f"{int(img_id):012d}.jpg")
+
+            img_tensor = load_image(img_path, preprocess).to(device)
+
+
+            pred_caption = generate_caption(model, tokenizer, img_tensor, device=device)
+            gt_caps = captions[str(img_id)]["captions"]
+
+            # Tokenized refs for BLEU
+            refs_tokenized.append([c.split() for c in gt_caps])
+
+            # String refs for JSON
+            refs_strings.append(gt_caps)
+
+            preds.append(pred_caption)
+
+            #if idx >= 20:
+            #    break
+
+        # Print 20 sample predictions
+        print("\nSample Predictions:\n")
+        num_examples = 20
+        for i in range(min(num_examples, len(preds))):
+            img_id = val_ids[i]
+            print(f"Image ID: {img_id}")
+            print(f"Prediction: {preds[i]}")
+            print(f"Ground Truths:")
+            for ref in refs_strings[i]:
+                print(f"  - {ref}")
+            print("-" * 60)
+
+
+        #print("Number of preds:", len(preds))
+        #print("Number of refs_tokenized:", len(refs_tokenized))
+        #print("Example hypothesis:", preds[0])
+        #print("Example hypothesis tokens:", preds[0].split())
+        #print("Example references:", refs_strings[0])
+        #print("Example references tokenized:", refs_tokenized[0])
+
+        #if HAS_COCOEVAL:
+            # Show first 2 examples only
+        #    for i in range(min(2, len(preds))):
+        #        img_id = str(int(val_ids[i]))
+        #        print(f"\nImage ID: {img_id}")
+        #        print("  COCOEvalCap refs (list of strings):")
+        #        print(" ", captions[img_id]["captions"])
+        #        print("  COCOEvalCap pred:")
+        #        print(" ", preds[i])
+
+            
+        # BLEU
+        smoothie = SmoothingFunction().method3
+
+        bleu1 = corpus_bleu(
+            refs_tokenized, [p.split() for p in preds],
+            weights=(1, 0, 0, 0),
+            smoothing_function=smoothie
+        )
+
+        bleu4 = corpus_bleu(
+            refs_tokenized, [p.split() for p in preds],
+            weights=(0.25, 0.25, 0.25, 0.25),
+            smoothing_function=smoothie
+        )
+
+        scores = {"BLEU-1": bleu1, "BLEU-4": bleu4}
+
+        # CIDEr / ROUGE
+        if HAS_COCOEVAL:
+            cider_refs = {}
+            cider_preds = {}
+
+            for i in range(len(preds)):
+                img_id = val_ids[i]
+                cid = str(int(img_id))  
+                cider_refs[cid] = captions[cid]["captions"]
+                cider_preds[cid] = [preds[i]]
+
+            #keys = list(cider_refs.keys())[:5]
+            #for k in keys:
+            #    print(f"{k}: {cider_refs[k]}")
+
+            #keys = list(cider_preds.keys())[:5]
+            #for k in keys:
+            #    print(f"{k}: {cider_preds[k]}")
+        
+            cider = Cider()
+            cider_score, _ = cider.compute_score(cider_refs, cider_preds)
+            scores["CIDEr"] = cider_score
+
+            rouge = Rouge()
+            rouge_score, _ = rouge.compute_score(cider_refs, cider_preds)
+            scores["ROUGE-L"] = rouge_score
+
+        # Save all samples
+        samples_full = []
+        for i in range(len(preds)):
+            img_id = val_ids[i]
+            samples_full.append({
+                "id": int(img_id),
+                "prediction": preds[i],
+                "references": refs_strings[i],
+                "image": f"{int(img_id):012d}.jpg",
+            })
+
+        # Save a preview subset (first 20)
+        samples_preview = samples_full[:20]
+        
+        param_info = count_encoder_decoder_params(model)
+
+        out_path = os.path.join(save_dir, "eval_results.json")
+
+        with open(out_path, "w") as f:
+            json.dump({
+                "scores": scores,
+                "derived_params": param_info,
+                "samples_preview": samples_preview,
+                "samples_full": samples_full
+            }, f, indent=2)
+
+        # Print final scores
+        print("\nEvaluation Scores:")
+        for k, v in scores.items():
+            print(f"{k}: {v:.4f}")
+
+        print(f"\nSaved detailed results to: {out_path}")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--checkpoint", type=str, default="checkpoints/vision_t5")
+    parser.add_argument("--data_dir", type=str, default="data/processed")
+
+    args = parser.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Device: {device}")
+
+    model, tokenizer, meta, config = load_experiment(args.checkpoint, device=device)
+    image_size = config["model"].get("image_size", 224)
+    preprocess = build_coco_transform(image_size=image_size)
+
+    evaluate(
+        model,
+        tokenizer,
+        preprocess=preprocess,
+        data_dir=args.data_dir,
+        save_dir=args.checkpoint,
+        device=device
+    )
+    
+    # python evaluate.py --checkpoint checkpoints/vision_t5/20251117_171912

src/evaluate_batched.py

@@ -0,0 +1,199 @@
+# evaluate_batched.py
+import os
+import json
+from tqdm import tqdm
+import torch
+from nltk.translate.bleu_score import corpus_bleu, SmoothingFunction
+from src.utils import count_encoder_decoder_params, load_experiment
+from src.inference import load_image, generate_caption
+from data.transforms import build_coco_transform
+
+try:
+    from pycocoevalcap.cider.cider import Cider
+    from pycocoevalcap.rouge.rouge import Rouge
+    HAS_COCOEVAL = True
+except ImportError:
+    print("WARNING: pycocoevalcap not installed → CIDEr/ROUGE disabled.")
+    HAS_COCOEVAL = False
+
+# Batched Evaluation (non-breaking addition)
+
+@torch.no_grad()
+def evaluate_batched(
+    model,
+    tokenizer,
+    preprocess,
+    image_size,
+    data_dir="data/processed",
+    save_dir="checkpoints",
+    device="cuda",
+    batch_size=16,
+    num_beams=1,
+):
+    """
+    Batched version of evaluate().
+    """
+
+    from src.inference import load_images_batch, generate_captions_batch
+
+    captions_path = os.path.join(data_dir, "captions.json")
+    splits_path = os.path.join(data_dir, "splits.json")
+
+    captions = json.load(open(captions_path))
+    splits = json.load(open(splits_path))
+    val_ids = splits["val"]
+
+    preds = []
+    refs_tokenized = []
+    refs_strings = []
+
+    print(f"Running *batched* evaluation on {len(val_ids)} images… (batch={batch_size})\n")
+
+    # Loop in batches
+    for start in tqdm(range(0, len(val_ids), batch_size), desc="Evaluating (batched)"):
+        end = min(start + batch_size, len(val_ids))
+        batch_ids = val_ids[start:end]
+
+        # Image paths
+        img_paths = [
+            os.path.join(data_dir, "images", f"{int(i):012d}.jpg")
+            for i in batch_ids
+        ]
+
+        # Load batch into tensor
+        img_batch = load_images_batch(img_paths, preprocess, image_size).to(device)
+
+        # Generate predictions for batch
+        batch_preds = generate_captions_batch(
+            model,
+            tokenizer,
+            img_batch,
+            device=device,
+            num_beams=num_beams,
+            max_new_tokens=32
+        )
+
+        # Collect references
+        for i, img_id in enumerate(batch_ids):
+            gt_caps = captions[str(img_id)]["captions"]
+
+            refs_strings.append(gt_caps)
+            refs_tokenized.append([c.split() for c in gt_caps])
+            preds.append(batch_preds[i])
+
+    # Print sample predictions (20 samples, same as evaluate())
+    print("\nSample Predictions:\n")
+    num_examples = 20
+    for i in range(min(num_examples, len(preds))):
+        img_id = val_ids[i]
+        print(f"Image ID: {img_id}")
+        print(f"Prediction: {preds[i]}")
+        print("Ground Truths:")
+        for ref in refs_strings[i]:
+            print(f"  - {ref}")
+        print("-" * 60)
+
+    # Compute BLEU
+    smoothie = SmoothingFunction().method3
+
+    bleu1 = corpus_bleu(
+        refs_tokenized, [p.split() for p in preds],
+        weights=(1, 0, 0, 0),
+        smoothing_function=smoothie
+    )
+
+    bleu4 = corpus_bleu(
+        refs_tokenized, [p.split() for p in preds],
+        weights=(0.25, 0.25, 0.25, 0.25),
+        smoothing_function=smoothie
+    )
+
+    scores = {"BLEU-1": bleu1, "BLEU-4": bleu4}
+
+    # CIDEr / ROUGE
+    if HAS_COCOEVAL:
+        cider_refs = {}
+        cider_preds = {}
+
+        for i in range(len(preds)):
+            img_id = val_ids[i]
+            cid = str(int(img_id))
+            cider_refs[cid] = captions[cid]["captions"]
+            cider_preds[cid] = [preds[i]]
+
+        cider = Cider()
+        cider_score, _ = cider.compute_score(cider_refs, cider_preds)
+        scores["CIDEr"] = cider_score
+
+        rouge = Rouge()
+        rouge_score, _ = rouge.compute_score(cider_refs, cider_preds)
+        scores["ROUGE-L"] = rouge_score
+
+    # Save results
+    samples_full = []
+    for i in range(len(preds)):
+        img_id = val_ids[i]
+        samples_full.append({
+            "id": int(img_id),
+            "prediction": preds[i],
+            "references": refs_strings[i],
+            "image": f"{int(img_id):012d}.jpg",
+        })
+
+    samples_preview = samples_full[:20]
+    param_info = count_encoder_decoder_params(model)
+
+    out_path = os.path.join(save_dir, "eval_results.json")
+    with open(out_path, "w") as f:
+        json.dump({
+            "scores": scores,
+            "derived_params": param_info,
+            "samples_preview": samples_preview,
+            "samples_full": samples_full
+        }, f, indent=2)
+
+    # Print final scores
+    print("\nEvaluation Scores:")
+    for k, v in scores.items():
+        print(f"{k}: {v:.4f}")
+
+    print(f"\nSaved batched results to: {out_path}")
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--checkpoint", type=str, required=True,
+                        help="Path to checkpoint directory")
+    parser.add_argument("--data_dir", type=str, default="data/processed")
+    parser.add_argument("--batch_size", type=int, default=16,
+                        help="Batch size for batched evaluation")
+    parser.add_argument("--num_beams", type=int, default=1,
+                        help="For beam search")
+
+    args = parser.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Device: {device}")
+
+    # Load model + tokenizer
+    model, tokenizer, meta, config = load_experiment(args.checkpoint, device=device)
+    image_size = config["model"].get("image_size", 224)
+    preprocess = build_coco_transform(image_size=image_size)
+
+    # Run batched evaluation
+    evaluate_batched(
+        model=model,
+        tokenizer=tokenizer,
+        preprocess=preprocess,
+        image_size=image_size,
+        data_dir=args.data_dir,
+        save_dir=args.checkpoint,
+        device=device,
+        batch_size=args.batch_size,
+        num_beams=args.num_beams
+    )
+
+    # Usage:
+    # python evaluate_batched.py --checkpoint checkpoints/vision_t5/20251117_171912 --batch_size 16

src/inference.py

@@ -0,0 +1,151 @@
+# inference.py
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from torchvision import transforms
+
+from transformers import T5TokenizerFast
+from transformers.modeling_outputs import BaseModelOutput
+from models.vision_t5 import VisionT5
+from src.utils import load_experiment
+from data.transforms import build_coco_transform
+
+
+
+def load_image(path, preprocess):
+    img = Image.open(path).convert("RGB")
+    return preprocess(img).unsqueeze(0)   # (1, 3, H, W)
+
+
+
+@torch.no_grad()
+def generate_caption(model, tokenizer, image_tensor, max_new_tokens=32, num_beams=1, device=None):
+    if device is None:
+        device = next(model.parameters()).device
+
+    model.eval()
+
+    image_tensor = image_tensor.to(device)
+
+    # Encode image
+    vision_out = model.vision_encoder(image_tensor)
+    img_embeds = vision_out["image_embeds"]
+
+    if img_embeds.dim() == 2:
+        img_embeds = img_embeds.unsqueeze(1)
+
+    projected = model.projector(img_embeds)
+
+    encoder_outputs = BaseModelOutput(last_hidden_state=projected)
+
+    start_token = model.t5.config.decoder_start_token_id
+
+    # explicit decoder inputs & mask (FIXES THE ERROR)
+    input_ids = torch.tensor([[start_token]], device=device)
+    attention_mask = torch.tensor([[1]], device=device)
+
+    output_ids = model.t5.generate(
+        encoder_outputs=encoder_outputs,
+        decoder_start_token_id=start_token,
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        num_beams=num_beams,
+        max_new_tokens=max_new_tokens,
+    )
+
+    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return caption
+
+
+
+
+# Batched evaluation helpers (non-breaking)
+
+@torch.no_grad()
+def load_images_batch(paths, preprocess, image_size):
+    resize = transforms.Resize((image_size, image_size))
+    tensors = []
+    for p in paths:
+        img = Image.open(p).convert("RGB")
+        img = resize(img)
+        t = preprocess(img).unsqueeze(0)
+        tensors.append(t)
+    return torch.cat(tensors, dim=0)
+
+
+@torch.no_grad()
+def generate_captions_batch(
+    model,
+    tokenizer,
+    image_batch,          # (B, 3, H, W)
+    max_new_tokens=32,
+    num_beams=1,
+    device=None,
+):
+    """
+    Batched version of generate_caption().
+    Does NOT replace or modify existing generate_caption().
+    """
+
+    if device is None:
+        device = next(model.parameters()).device
+
+    model.eval()
+
+    image_batch = image_batch.to(device)
+
+    # Encode in batch
+    vision_out = model.vision_encoder(image_batch)
+    img_embeds = vision_out["image_embeds"]   # (B, D) or (B, S, D)
+
+    if img_embeds.dim() == 2:
+        img_embeds = img_embeds.unsqueeze(1)
+
+    projected = model.projector(img_embeds)    # (B, S, d_model)
+    encoder_outputs = BaseModelOutput(last_hidden_state=projected)
+
+    # Build batched decoder inputs
+    start = model.t5.config.decoder_start_token_id
+    B = image_batch.size(0)
+
+    input_ids = torch.full((B, 1), start, dtype=torch.long, device=device)
+    attention_mask = torch.ones((B, 1), dtype=torch.long, device=device)
+
+    # Standard HF batching
+    output_ids = model.t5.generate(
+        encoder_outputs=encoder_outputs,
+        decoder_start_token_id=start,
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        num_beams=num_beams,
+        max_new_tokens=max_new_tokens,
+    )
+
+    # Decode individually
+    return [
+        tokenizer.decode(ids, skip_special_tokens=True)
+        for ids in output_ids
+    ]
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--image", type=str, required=True, help="Path to image")
+    parser.add_argument("--checkpoint", type=str, default="checkpoints/vision_t5")
+
+    args = parser.parse_args()
+
+    # Load model + tokenizer + config
+    model, tokenizer, meta, config = load_experiment(args.checkpoint)
+    image_size = config["model"].get("image_size", 224)
+    preprocess = build_coco_transform(image_size)
+
+    # Load image
+    image_tensor = load_image(args.image, preprocess)
+
+    # Generate caption
+    caption = generate_caption(model, tokenizer, image_tensor)
+    print("\nCaption:", caption)
+

src/interpretability.py

@@ -0,0 +1,466 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+import matplotlib.pyplot as plt
+from torchvision.transforms.functional import resize
+from transformers.modeling_outputs import BaseModelOutput
+import cv2
+from transformers.models.vit.modeling_vit import ViTModel
+
+
+import torch
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+
+
+class GradCAM:
+    def __init__(self, vision_encoder):
+        self.model = vision_encoder.model
+        self.target_layer = self._find_last_conv_layer()
+
+        self.activations = None
+        self.gradients = None
+
+        self.target_layer.register_forward_hook(self._hook_forward)
+        self.target_layer.register_backward_hook(self._hook_backward)
+
+    def _find_last_conv_layer(self):
+        for module in reversed(list(self.model.modules())):
+            if isinstance(module, torch.nn.Conv2d):
+                return module
+        raise RuntimeError("No Conv2D layer found for Grad-CAM.")
+
+    def _hook_forward(self, module, inp, out):
+        self.activations = out.detach()
+
+    def _hook_backward(self, module, grad_in, grad_out):
+        self.gradients = grad_out[0].detach()
+
+    def generate(self, image_tensor):
+        self.model.zero_grad()
+
+        out = self.model(image_tensor)  # (B, C, H, W)
+
+        if out.ndim == 4:
+            pooled = out.mean(dim=[2, 3])  # (B, C)
+        elif out.ndim == 3:
+            pooled = out.mean(dim=1)
+        else:
+            pooled = out
+
+        score = pooled.norm()
+        score.backward()
+
+        weights = self.gradients.mean(dim=(2, 3), keepdim=True)
+        cam = (weights * self.activations).sum(dim=1).squeeze()
+
+        cam = F.relu(cam)
+        cam -= cam.min()
+        cam /= cam.max() + 1e-8
+
+        return cam.cpu().numpy()
+
+    def save(self, img_tensor, save_path):
+        cam = self.generate(img_tensor)
+
+        img_np = img_tensor[0].permute(1, 2, 0).cpu().numpy()
+        img_np = (img_np - img_np.min()) / (img_np.max() - img_np.min())
+
+        cam_resized = cv2.resize(cam, (img_np.shape[1], img_np.shape[0]))
+
+        plt.figure(figsize=(6, 6))
+        plt.imshow(img_np)
+        plt.imshow(cam_resized, cmap="inferno", alpha=0.45)
+        plt.axis("off")
+        plt.tight_layout()
+        plt.savefig(save_path, dpi=300, bbox_inches="tight")
+        plt.close()
+
+        print(f"[GradCAM] Saved to {save_path}")
+
+
+def get_vit_self_attention(model, image_tensor):
+    vision = model.vision_encoder
+
+    if "Resnet" in type(vision).__name__:
+        return None
+
+    # Check for CLIP
+    if hasattr(vision, "model"):
+
+        if hasattr(vision.model, "vision_model"):
+            hf_vit = vision.model.vision_model
+
+            outputs = hf_vit(
+                pixel_values=image_tensor,
+                output_attentions=True,
+                return_dict=True,
+            )
+            return outputs.attentions
+
+    # Check for ViT 
+    if isinstance(vision.model, ViTModel):
+
+        outputs = vision.model(
+            pixel_values=image_tensor,
+            output_attentions=True,
+            return_dict=True,
+        )
+        return outputs.attentions
+
+    raise ValueError("Vision encoder does not expose ViT attentions.")
+
+
+
+# ATTENTION ROLLOUT (across layers)
+def attention_rollout(attn_mats, discard_ratio=0.0):
+
+    device = attn_mats[0].device
+    result = torch.eye(attn_mats[0].size(-1), device=device)
+
+    for attn in attn_mats:
+        attn = attn.mean(dim=0)  # average heads
+
+        if discard_ratio > 0:
+            flat = attn.view(-1)
+            threshold = flat.topk(int(flat.numel() * discard_ratio), largest=False)[0].max()
+            attn = torch.where(attn < threshold, torch.zeros_like(attn), attn)
+
+        attn = attn / attn.sum(dim=-1, keepdim=True)
+        result = attn @ result
+
+    return result  
+
+
+def rollout_to_image(rollout, image_size):
+    tokens = rollout.size(0)
+    num_patches = int((tokens - 1) ** 0.5)
+
+    spatial = rollout[0, 1:].reshape(num_patches, num_patches)
+    spatial = (spatial - spatial.min()) / (spatial.max() - spatial.min())
+
+    spatial = resize(
+        spatial.unsqueeze(0).unsqueeze(0),
+        (image_size, image_size)
+    )
+
+    return spatial.squeeze().detach().cpu().numpy()
+
+
+def plot_attention_overlay(image, heatmap, alpha=0.45):
+    if torch.is_tensor(image):
+        image = image.permute(1,2,0).cpu().numpy()
+
+    image = (image - image.min()) / (image.max() - image.min())
+
+    plt.figure(figsize=(6,6))
+    plt.imshow(image)
+    plt.imshow(heatmap, cmap='inferno', alpha=alpha)
+    plt.axis("off")
+    plt.show()
+
+
+# GRADIENT MAP
+def token_gradient_map(model, tokenizer, image_tensor, target_word, device="cuda"):
+    model.eval()
+
+    image_tensor = image_tensor.to(device)
+    image_tensor.requires_grad_(True)
+
+    vision_out = model.vision_encoder(image_tensor)
+    img_embeds = vision_out["image_embeds"]
+
+    if img_embeds.dim() == 2:
+        img_embeds = img_embeds.unsqueeze(1)
+
+    projected = model.projector(img_embeds)
+
+    encoder_outputs = BaseModelOutput(last_hidden_state=projected)
+
+    start = model.t5.config.decoder_start_token_id
+    decoder_input_ids = torch.tensor([[start]], device=device)
+
+    outputs = model.t5(
+        encoder_outputs=encoder_outputs,
+        decoder_input_ids=decoder_input_ids,
+        return_dict=True,
+    )
+
+    logits = outputs.logits[:, -1, :]
+    target_id = tokenizer.convert_tokens_to_ids(target_word)
+    logit = logits[0, target_id]
+
+    logit.backward()
+
+    grad = image_tensor.grad.abs().mean(dim=1).squeeze().cpu().numpy()
+    grad = (grad - grad.min()) / (grad.max() - grad.min() + 1e-8)
+
+    return grad
+
+# ATTENTION x GRAD
+def attngrad(model, tokenizer, image_tensor, target_word, image_size=224, device="cuda"):
+
+    raw_attns = get_vit_self_attention(model, image_tensor.to(device))
+    attn_mats = [a[0] for a in raw_attns]
+
+    rollout = attention_rollout(attn_mats)
+    roll_map = rollout_to_image(rollout, image_size)
+
+    grad_map = token_gradient_map(model, tokenizer, image_tensor, target_word, device)
+
+    combined = roll_map * grad_map
+    combined = (combined - combined.min()) / (combined.max() - combined.min())
+    return combined
+
+
+
+def token_gradient_map_smooth(model, tokenizer, image_tensor, target_word, sigma=5, device="cuda"):
+    model.eval()
+
+    image_tensor = image_tensor.to(device)
+    image_tensor.requires_grad_(True)
+
+    # Vision encoder
+    vision_out = model.vision_encoder(image_tensor)
+    img_embeds = vision_out["image_embeds"]
+    if img_embeds.dim() == 2:
+        img_embeds = img_embeds.unsqueeze(1)
+
+    projected = model.projector(img_embeds)
+
+
+    encoder_outputs = BaseModelOutput(last_hidden_state=projected)
+
+
+    start_token = model.t5.config.decoder_start_token_id
+
+    decoder_input_ids = torch.tensor(
+        [[start_token]], device=device, dtype=torch.long
+    )
+    attention_mask = torch.tensor([[1]], device=device)
+
+    outputs = model.t5(
+        encoder_outputs=encoder_outputs,
+        decoder_input_ids=decoder_input_ids,
+        attention_mask=attention_mask,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=True,
+    )
+
+    vocab_logits = outputs.logits[:, -1, :]
+    target_id = tokenizer.convert_tokens_to_ids(target_word)
+    logit = vocab_logits[0, target_id]
+
+    logit.backward()
+
+    grad = image_tensor.grad.data.abs().mean(dim=1).squeeze().cpu().numpy()
+    grad = (grad - grad.min()) / (grad.max() - grad.min() + 1e-8)
+
+    grad_smooth = smooth_heatmap(grad, sigma=sigma)
+    return grad_smooth
+
+
+def integrated_gradients(
+    model,
+    tokenizer,
+    image_tensor,
+    target_word,
+    steps=30,
+    device="cuda"
+):
+    model.eval()
+    device = torch.device(device)
+
+    image_tensor = image_tensor.to(device)
+    image_tensor.requires_grad_(True)
+
+    baseline = torch.zeros_like(image_tensor)
+
+    target_id = tokenizer.convert_tokens_to_ids(target_word)
+
+    total_grad = torch.zeros_like(image_tensor)
+
+    for i in range(1, steps + 1):
+        alpha = i / steps
+
+        img = baseline + alpha * (image_tensor - baseline)
+        img.requires_grad_(True)
+
+        vision_out = model.vision_encoder(img)
+        img_embeds = vision_out["image_embeds"]
+        if img_embeds.dim() == 2:
+            img_embeds = img_embeds.unsqueeze(1)
+
+        projected = model.projector(img_embeds)
+        encoder_outputs = BaseModelOutput(last_hidden_state=projected)
+
+        start_token = model.t5.config.decoder_start_token_id
+        decoder_input_ids = torch.tensor([[start_token]], device=device)
+        attention_mask = torch.tensor([[1]], device=device)
+
+        outputs = model.t5(
+            encoder_outputs=encoder_outputs,
+            decoder_input_ids=decoder_input_ids,
+            attention_mask=attention_mask, 
+            return_dict=True,
+        )
+
+        vocab_logits = outputs.logits[:, -1, :]
+        logit = vocab_logits[0, target_id]
+
+        grads = torch.autograd.grad(
+            outputs=logit,
+            inputs=img,
+            retain_graph=True,
+            create_graph=False,
+            allow_unused=True,  
+        )[0]
+
+        if grads is None:
+            raise RuntimeError("Integrated gradients: grad is None — gradient path was broken.")
+
+        total_grad += grads
+
+    avg_grad = total_grad / steps
+
+    heat = avg_grad.abs().mean(dim=1).squeeze().cpu().numpy()
+    heat = (heat - heat.min()) / (heat.max() - heat.min() + 1e-8)
+
+    return heat
+
+
+
+def smooth_heatmap(hm, k=21, sigma=6):
+    hm = cv2.GaussianBlur(hm, (k, k), sigma)
+    hm = (hm - hm.min()) / (hm.max() - hm.min() + 1e-8)
+    return hm
+
+
+def get_cross_attention(model, encoder_outputs, decoder_input_ids, device="cuda"):
+
+    model.eval()
+    with torch.no_grad():
+        outputs = model.t5(
+            encoder_outputs=encoder_outputs,
+            decoder_input_ids=decoder_input_ids.to(device),
+            output_attentions=True,
+            return_dict=True,
+        )
+
+    # outputs.cross_attentions is a tuple of layers (batch, heads, tgt_len, src_len)
+    cross = outputs.cross_attentions
+    attn_layers = [c[0] for c in cross]  # use batch 0
+    return attn_layers
+
+
+"""
+def cross_attention_to_image(attn, image_size=224):
+
+    attn = attn.mean(dim=0)   # (tgt_len, src_len)
+
+    attn = attn[-1]           # (src_len,)
+
+    attn = attn[1:]
+
+    num_patches = int(attn.numel() ** 0.5)
+    heat = attn.reshape(num_patches, num_patches)
+
+    heat = heat - heat.min()
+    heat = heat / (heat.max() + 1e-8)
+
+    heat = resize(
+        heat.unsqueeze(0).unsqueeze(0),
+        (image_size, image_size)
+    ).squeeze()
+
+    return heat.detach().cpu().numpy()
+"""
+
+def cross_attention_to_image(attn):
+
+    attn = torch.tensor(attn) if not torch.is_tensor(attn) else attn
+
+    if attn.numel() == 0:
+        return np.zeros((14, 14), dtype=np.float32)
+
+    if attn.dim() == 2:
+        attn_vec = attn[-1]  # use last generated token
+    elif attn.dim() == 1:
+        attn_vec = attn
+    else:
+        raise ValueError(f"Unexpected attn shape: {attn.shape}")
+
+    # DROP CLS TOKEN (index 0) for CLIP ViT-L/14 197 tokens but 196 spatial patches
+    if attn_vec.size(0) == 197:
+        attn_vec = attn_vec[1:]      # now length = 196
+
+    src_len = attn_vec.size(0)
+    side = int(src_len**0.5)
+
+    if side * side != src_len:
+        new_len = side * side
+        padded = torch.zeros(new_len, device=attn_vec.device)
+        padded[:min(new_len, src_len)] = attn_vec[:min(new_len, src_len)]
+        attn_vec = padded
+
+    attn_vec = attn_vec / (attn_vec.max() + 1e-8)
+
+    heatmap = attn_vec.reshape(side, side).cpu().numpy()
+
+    return heatmap
+
+
+
+def plot_cross_attention_overlay(image_tensor, heatmap, save_path=None, alpha=0.45):
+    img = image_tensor[0].permute(1,2,0).cpu().numpy()
+    img = (img - img.min()) / (img.max() - img.min())
+
+    plt.figure(figsize=(6,6))
+    plt.imshow(img)
+    plt.imshow(heatmap, cmap='inferno', alpha=alpha)
+    plt.axis("off")
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches="tight")
+        plt.close()
+        print(f"[CrossAttention] Saved to {save_path}")
+    else:
+        plt.show()
+
+
+def visualize_cross_attention(model, tokenizer, image_tensor, word, device="cuda"):
+    device = torch.device(device)
+    image_tensor = image_tensor.to(device)
+
+    vision_out = model.vision_encoder(image_tensor)
+    img_embeds = vision_out["image_embeds"]
+    if img_embeds.dim() == 2:
+        img_embeds = img_embeds.unsqueeze(1)
+
+    projected = model.projector(img_embeds)
+    encoder_outputs = BaseModelOutput(last_hidden_state=projected)
+
+    generated = [model.t5.config.decoder_start_token_id]
+
+    for _ in range(30):
+        decoder_input_ids = torch.tensor([generated], device=device)
+        attn_layers = get_cross_attention(
+            model, encoder_outputs, decoder_input_ids
+        )
+
+        logits = model.t5(
+            encoder_outputs=encoder_outputs,
+            decoder_input_ids=decoder_input_ids,
+            return_dict=True
+        ).logits[:, -1, :]
+        next_id = int(logits.argmax())
+        generated.append(next_id)
+
+        if next_id == tokenizer.convert_tokens_to_ids(word):
+            break
+
+    last_attn = attn_layers[-1]  # (heads, T, S)
+    heat = cross_attention_to_image(last_attn)
+
+    plot_cross_attention_overlay(image_tensor, heat)

src/train.py

@@ -0,0 +1,176 @@
+import argparse
+import yaml
+import torch
+from torch.optim import AdamW
+from tqdm import tqdm
+from transformers import T5TokenizerFast
+
+
+from models.vision_t5 import VisionT5
+from models.encoder_projection_t5 import ImageProjection
+import models.encoders as encoders
+
+from data.loaders import get_coco_dataloaders
+from src.inference import generate_caption
+from src.utils import save_experiment, filter_kwargs, build_model
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+import math
+
+
+
+def build_cosine_warmup_scheduler(optimizer, num_warmup_steps, num_training_steps):
+    def lr_lambda(step):
+        if step < num_warmup_steps:
+            return float(step) / float(max(1, num_warmup_steps))
+        progress = float(step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+        return 0.5 * (1 + math.cos(math.pi * progress))  # cosine decay
+
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+
+def train_one_epoch(model, dataloader, optimizer, device, scaler, scheduler):
+    model.train()
+    running_loss = 0.0
+
+    for batch in tqdm(dataloader, desc="Training"):
+        pixel_values = batch["pixel_values"].to(device)
+        input_ids = batch["input_ids"].to(device)
+        attention_mask = batch["attention_mask"].to(device)
+
+        # teacher forcing labels
+        labels = input_ids.clone()
+        labels[labels == model.t5.config.pad_token_id] = -100 # HF provided value to ignore in labels for loss calc.
+
+        optimizer.zero_grad()
+
+        # Using AMP to save memory
+        with torch.cuda.amp.autocast():
+            outputs = model(
+                pixel_values=pixel_values,
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                labels=labels,
+            )
+            loss = outputs.loss
+
+        scaler.scale(loss).backward()
+
+        # Gradient clipping
+        scaler.unscale_(optimizer) 
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
+        scaler.step(optimizer)
+        scaler.update()
+        scheduler.step()
+        running_loss += loss.item()
+
+    return running_loss / len(dataloader)
+
+
+# Validation
+@torch.no_grad()
+def validate(model, tokenizer, dataloader, device, preview=False):
+    model.eval()
+    running_loss = 0.0
+
+    sample_img = None
+    sample_gt = None
+
+    for batch in tqdm(dataloader, desc="Validation"):
+        pixel_values = batch["pixel_values"].to(device)
+        input_ids = batch["input_ids"].to(device)
+        attention_mask = batch["attention_mask"].to(device)
+
+        # Teacher-forcing labels
+        labels = input_ids.clone()
+        labels[labels == tokenizer.pad_token_id] = -100
+
+        outputs = model(
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels,
+        )
+
+        running_loss += outputs.loss.item()
+
+        # Store sample for preview
+        if preview and sample_img is None:
+            sample_img = pixel_values[0].detach().cpu()
+            # decode GT caption (first non-pad tokens)
+            gt_ids = input_ids[0][input_ids[0] != tokenizer.pad_token_id]
+            sample_gt = tokenizer.decode(gt_ids, skip_special_tokens=True)
+
+    # preview
+    if preview and sample_img is not None:
+        print("\n--- Validation Preview ---")
+        pred = generate_caption(model, tokenizer, sample_img.unsqueeze(0), device=device)
+        print("Prediction:", pred)
+        print("Ground Truth:", sample_gt)
+        print("--------------------------\n")
+
+    return running_loss / len(dataloader)
+
+
+
+def main(config):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Model + tokenizer
+    model, tokenizer = build_model(config)
+    model.to(device)
+
+    # Data
+    batch_size = config["training"]["batch_size"]
+    image_size = config["model"].get("image_size", 224)
+    train_loader, val_loader, _ = get_coco_dataloaders(batch_size=batch_size, data_dir=config["paths"]["data_dir"], image_size=image_size)
+
+    optimizer = AdamW(model.parameters(), lr=config["training"]["lr"])
+    scaler = torch.cuda.amp.GradScaler() # For mixed precision
+
+    num_training_steps = len(train_loader) * config["training"]["epochs"]
+    num_warmup_steps = int(0.05 * num_training_steps)
+    scheduler = build_cosine_warmup_scheduler(
+        optimizer,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps
+    )
+    best_val = float("inf")
+    best_epoch = -1
+
+    # Train loop
+    for epoch in range(1, config["training"]["epochs"] + 1):
+        print(f"\nEpoch {epoch}/{config['training']['epochs']}")
+
+        train_loss = train_one_epoch(model, train_loader, optimizer, device, scaler, scheduler)
+        print("Train Loss:", train_loss)
+
+        val_loss = validate(model, tokenizer, val_loader, device, preview=config["training"]["preview_val"])
+        print("Val Loss:", val_loss)
+
+        if val_loss < best_val:
+            best_val = val_loss
+            best_epoch = epoch
+
+            save_experiment(
+                model=model,
+                tokenizer=tokenizer,
+                config=config,
+                save_dir=config["paths"]["output_dir"],
+                notes=f"BEST checkpoint epoch={epoch}, val_loss={val_loss:.4f}"
+            )
+            print(f"[CHECKPOINT] Saved new BEST model at epoch {epoch}")
+            
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, required=True)
+    args = parser.parse_args()
+
+    with open(args.config, "r") as f:
+        config = yaml.safe_load(f)
+
+    main(config)

src/utils.py

@@ -0,0 +1,256 @@
+# utils.py
+import os
+import yaml
+import torch
+from datetime import datetime
+
+from transformers import T5TokenizerFast
+from models.vision_t5 import VisionT5
+import models.encoders as encoders
+from models.encoder_projection_t5 import ImageProjection
+import inspect
+
+
+
+def timestamp():
+    return datetime.now().strftime("%Y%m%d_%H%M%S")
+
+
+
+def save_experiment(model, tokenizer, config, save_dir, notes="", run_name=None, add_timestamp=True):
+
+    if add_timestamp:
+        tag = timestamp()
+        if run_name:
+            save_dir = os.path.join(save_dir, f"{run_name}_{tag}")
+        else:
+            save_dir = os.path.join(save_dir, tag)
+
+    os.makedirs(save_dir, exist_ok=True)
+
+    torch.save(model.state_dict(), os.path.join(save_dir, "pytorch_model.bin"))
+
+    tok_dir = os.path.join(save_dir, "tokenizer")
+    os.makedirs(tok_dir, exist_ok=True)
+    tokenizer.save_pretrained(tok_dir)
+
+    with open(os.path.join(save_dir, "config_trained.yaml"), "w") as f:
+        yaml.safe_dump(config, f)
+
+    metadata = {
+        "encoder": config["model"]["encoder"],
+        "encoder_params": config["model"].get("encoder_params", {}),
+        "decoder": config["model"]["t5_name"],
+        "decoder_params": config["model"].get("decoder_params", {}),
+        "train_epochs": config["training"]["epochs"],
+        "batch_size": config["training"]["batch_size"],
+        "lr": config["training"]["lr"],
+        "notes": notes,
+        "run_name": run_name,
+        "timestamp": timestamp(),
+    }
+
+    with open(os.path.join(save_dir, "metadata.yaml"), "w") as f:
+        yaml.safe_dump(metadata, f)
+
+    print(f"[OK] Experiment saved → {save_dir}")
+    return save_dir
+
+
+
+
+def load_experiment(checkpoint_dir, device="cpu"):
+    import yaml, torch, os
+
+    metadata_path = os.path.join(checkpoint_dir, "metadata.yaml")
+    config_path = os.path.join(checkpoint_dir, "config_trained.yaml")
+
+    if not os.path.exists(metadata_path):
+        raise FileNotFoundError(f"No metadata.yaml found at {checkpoint_dir}")
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"No config_trained.yaml found at {checkpoint_dir}")
+
+    with open(metadata_path, "r") as f:
+        metadata = yaml.safe_load(f)
+
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+
+
+    model, tokenizer = build_model(config)
+
+    tok_dir = os.path.join(checkpoint_dir, "tokenizer")
+    if os.path.isdir(tok_dir):
+        tokenizer = T5TokenizerFast.from_pretrained(tok_dir)
+
+    ckpt_path = os.path.join(checkpoint_dir, "pytorch_model.bin")
+    weights = torch.load(ckpt_path, map_location=device)
+    model.load_state_dict(weights, strict=False)
+
+    model.to(device)
+    model.eval()
+
+    print(f"Loaded experiment from {checkpoint_dir}")
+    return model, tokenizer, metadata, config
+
+
+
+def filter_kwargs(cls, kwargs):
+    sig = inspect.signature(cls.__init__).parameters
+    return {k: v for k, v in kwargs.items() if k in sig}
+
+
+
+def build_model(config):
+
+    encoder_name = config["model"]["encoder"]
+    raw_encoder_params = config["model"].get("encoder_params", {})
+
+    t5_name = config["model"]["t5_name"]
+    decoder_params = config["model"].get("decoder_params", {})
+
+    tokenizer = T5TokenizerFast.from_pretrained(t5_name)
+
+    # dynamically load encoder class
+    if not hasattr(encoders, encoder_name):
+        raise ValueError(f"Encoder '{encoder_name}' not found in encoders.py")
+
+    EncoderClass = getattr(encoders, encoder_name)
+
+    encoder_params = filter_kwargs(EncoderClass, raw_encoder_params)
+
+    # Instantiate encoder
+    vision_encoder = EncoderClass(**encoder_params)
+
+    # Projection layer
+    t5_hidden = VisionT5.get_t5_hidden_size(t5_name)
+    projector = ImageProjection(
+        encoder_dim=vision_encoder.get_output_dim(),
+        t5_hidden_size=t5_hidden
+    )
+
+    # Construct model
+    model = VisionT5(
+        vision_encoder=vision_encoder,
+        projector=projector,
+        t5_name=t5_name,
+        decoder_params=decoder_params
+    )
+
+    return model, tokenizer
+
+
+
+def load_yaml(path):
+    with open(path, "r") as f:
+        return yaml.safe_load(f)
+
+
+
+def count_encoder_decoder_params(model):
+
+    enc_total = enc_train = 0
+    proj_total = proj_train = 0
+    dec_total = dec_train = 0
+    other_total = other_train = 0
+
+    for name, p in model.named_parameters():
+        n = p.numel()
+
+        # Vision Encoder
+        if name.startswith("vision_encoder."):
+            enc_total += n
+            if p.requires_grad:
+                enc_train += n
+            continue
+
+        # Projector
+        if name.startswith("projector."):
+            proj_total += n
+            if p.requires_grad:
+                proj_train += n
+            continue
+
+        # T5 Decoder (covers small, base, large, AND LoRA)
+        if (
+            name.startswith("t5.decoder.") or
+            "decoder.block" in name or
+            name.startswith("t5.model.decoder.") or
+            name.startswith("t5.lm_head.") or
+            name.startswith("t5.shared.")
+        ):
+            dec_total += n
+            if p.requires_grad:
+                dec_train += n
+            continue
+
+        if "lora_" in name and "decoder" in name:
+            dec_total += n
+            if p.requires_grad:
+                dec_train += n
+            continue
+
+        # T5 Encoder (always frozen)
+        if name.startswith("t5.encoder."):
+            other_total += n
+            if p.requires_grad:
+                other_train += n
+            continue
+
+        # Other params
+        other_total += n
+        if p.requires_grad:
+            other_train += n
+
+    total_params = enc_total + proj_total + dec_total + other_total
+    trainable_params = enc_train + proj_train + dec_train + other_train
+
+    return {
+        "encoder_total_params": enc_total,
+        "encoder_trainable_params": enc_train,
+        "encoder_trainable_fraction":
+            enc_train / enc_total if enc_total else None,
+
+        "projector_total_params": proj_total,
+        "projector_trainable_params": proj_train,
+        "projector_trainable_fraction":
+            proj_train / proj_total if proj_total else None,
+
+        "decoder_total_params": dec_total,
+        "decoder_trainable_params": dec_train,
+        "decoder_trainable_fraction":
+            dec_train / dec_total if dec_total else None,
+
+        "other_total_params": other_total,
+        "other_trainable_params": other_train,
+
+        "total_params": total_params,
+        "trainable_params": trainable_params,
+        "trainable_params_fraction":
+            trainable_params / total_params if total_params else None,
+    }
+
+
+def classify_param(name):
+
+  if name.startswith("vision_encoder."):
+      return "encoder"
+
+  if name.startswith("projector."):
+      return "projector"
+
+  if (
+      name.startswith("t5.decoder.") or
+      name.startswith("t5.model.decoder.") or
+      "decoder.block" in name or
+      name.startswith("t5.lm_head.") or
+      name.startswith("t5.shared.") or
+      ("lora_" in name and "decoder" in name)
+  ):
+      return "decoder"
+
+  if name.startswith("t5.encoder."):
+      return "t5_encoder_frozen"
+
+  return "other"
+