Update RemoteCLIP model package

conf/config.yaml

@@ -1,25 +1,36 @@
 seed: 42
 data:
-  path: data/remoteclip.npz
-  protocol: synthetic_remoteclip
-  train_samples: 32
-  test_samples: 8
-  image_size: 32
-  vocabulary_size: 64
-  context_length: 8
-  num_classes: 4
+  root: data
+  protocol: remoteclip_rgb_bpe_pairs_v1
+  image_size: 224
+  channels: 3
+  context_length: 77
+  vocabulary_size: 49408
+  sot_token_id: 49406
+  eot_token_id: 49407
+  pad_token_id: 0
+  train_samples: 8
+  test_samples: 4
+  num_semantic_groups: 2
 model:
-  embed_dim: 32
-  vision_width: 32
-  text_width: 32
+  image_size: 224
+  patch_size: 32
+  embed_dim: 64
+  vision_width: 64
+  vision_layers: 2
+  vision_heads: 4
+  text_width: 64
   text_layers: 2
   text_heads: 4
-train:
-  epochs: 5
-  batch_size: 8
-  learning_rate: 0.001
-  weight_decay: 0.0001
+training:
+  epochs: 1
+  batch_size: 4
+  learning_rate: 0.0005
+  weight_decay: 0.01
   num_workers: 0
+runtime:
+  device: auto
+  amp: true
 paths:
   checkpoint: result/checkpoints/remoteclip.pt
   training_metrics: result/training/metrics.json

config.json

@@ -5,25 +5,29 @@
     "RemoteCLIP"
   ],
   "framework": "PyTorch",
-  "domain": "earth-science",
+  "domain": "earth-observation",
   "task": "remote-sensing-image-text-retrieval",
   "implementation": {
     "entry_point": "model/remoteclip.py",
-    "scope": "compact image-text contrastive RemoteCLIP reproduction"
+    "scope": "CLIP-compatible vision and text dual encoders with multi-positive symmetric contrastive learning"
   },
   "architecture": {
-    "family": "dual-encoder vision-language model",
-    "input_format": "image-text NPZ pairs",
-    "image_channels": 3,
-    "image_size": 32,
-    "vocabulary_size": 64,
-    "context_length": 8,
-    "embed_dim": 32,
-    "vision_width": 32,
-    "text_width": 32,
+    "family": "vision transformer and causal text transformer dual encoder",
+    "image_shape": [
+      3,
+      224,
+      224
+    ],
+    "patch_size": 32,
+    "context_length": 77,
+    "vocabulary_size": 49408,
+    "embed_dim": 64,
+    "vision_width": 64,
+    "vision_layers": 2,
+    "vision_heads": 4,
+    "text_width": 64,
     "text_layers": 2,
-    "text_heads": 4,
-    "training_objective": "bidirectional InfoNCE"
+    "text_heads": 4
   },
   "data": {
     "datasets": [
@@ -31,19 +35,13 @@
       "RSICD",
       "UCM-Captions"
     ],
-    "official_repo": "ChenDelong1999/RemoteCLIP",
-    "protocol": "synthetic_remoteclip",
-    "synthetic_train_samples": 32,
-    "synthetic_test_samples": 8
+    "protocol": "remoteclip_rgb_bpe_pairs_v1",
+    "input_format": "image_text_npz",
+    "train_samples": 8,
+    "test_samples": 4
   },
-  "metrics": [
-    "image_to_text_r1",
-    "image_to_text_r5",
-    "text_to_image_r1",
-    "text_to_image_r5",
-    "mean_recall"
-  ],
   "configuration_sources": [
+    "configuration.json",
     "conf/config.yaml",
     "model/remoteclip.py"
   ]

configuration.json

@@ -2,8 +2,9 @@
   "framework": "PyTorch",
   "task": "remote_sensing_image_text_retrieval",
   "model": "RemoteCLIP",
+  "license": "Apache-2.0",
   "input_format": "image_text_npz",
-  "protocol": "synthetic_remoteclip",
+  "protocol": "remoteclip_rgb_bpe_pairs_v1",
   "default_config": "conf/config.yaml",
   "train": "scripts/train.py",
   "inference": "scripts/inference.py",

model/remoteclip.py

@@ -1,4 +1,4 @@
-"""Compact RemoteCLIP image-text contrastive model."""
+"""RemoteCLIP dual encoder with ViT and causal text Transformer."""
 
 import math
 
@@ -7,66 +7,100 @@ from torch import nn
 from torch.nn import functional as F
 
 
+class VisionTransformer(nn.Module):
+    def __init__(self, image_size, patch_size, width, layers, heads, output_dim):
+        super().__init__()
+        if image_size % patch_size:
+            raise ValueError("image_size must be divisible by patch_size")
+        patches = (image_size // patch_size) ** 2
+        self.patch_embed = nn.Conv2d(3, width, patch_size, patch_size, bias=False)
+        self.class_embedding = nn.Parameter(torch.empty(1, 1, width))
+        self.position_embedding = nn.Parameter(torch.empty(1, patches + 1, width))
+        layer = nn.TransformerEncoderLayer(
+            width, heads, width * 4, activation="gelu", batch_first=True,
+            norm_first=True, dropout=0.0,
+        )
+        self.transformer = nn.TransformerEncoder(layer, layers)
+        self.norm = nn.LayerNorm(width)
+        self.projection = nn.Parameter(torch.empty(width, output_dim))
+        nn.init.normal_(self.class_embedding, std=width ** -0.5)
+        nn.init.normal_(self.position_embedding, std=width ** -0.5)
+        nn.init.normal_(self.projection, std=width ** -0.5)
+
+    def forward(self, images):
+        tokens = self.patch_embed(images).flatten(2).transpose(1, 2)
+        cls = self.class_embedding.expand(images.shape[0], -1, -1)
+        tokens = torch.cat((cls, tokens), dim=1) + self.position_embedding
+        return self.norm(self.transformer(tokens)[:, 0]) @ self.projection
+
+
 class RemoteCLIP(nn.Module):
+    """CLIP-compatible encoders; EOT is the largest token id in each sequence."""
+
     def __init__(
         self,
-        vocabulary_size=64,
-        context_length=8,
-        embed_dim=32,
-        vision_width=32,
-        text_width=32,
+        vocabulary_size=49408,
+        context_length=77,
+        eot_token_id=49407,
+        image_size=224,
+        patch_size=32,
+        embed_dim=64,
+        vision_width=64,
+        vision_layers=2,
+        vision_heads=4,
+        text_width=64,
         text_layers=2,
         text_heads=4,
     ):
         super().__init__()
         self.context_length = context_length
-        self.image_encoder = nn.Sequential(
-            nn.Conv2d(3, vision_width, 5, stride=2, padding=2),
-            nn.GELU(),
-            nn.Conv2d(vision_width, vision_width * 2, 3, stride=2, padding=1),
-            nn.GELU(),
-            nn.AdaptiveAvgPool2d(1),
-            nn.Flatten(),
-            nn.Linear(vision_width * 2, embed_dim),
+        self.eot_token_id = eot_token_id
+        self.visual = VisionTransformer(
+            image_size, patch_size, vision_width, vision_layers, vision_heads, embed_dim
         )
         self.token_embedding = nn.Embedding(vocabulary_size, text_width, padding_idx=0)
-        self.position_embedding = nn.Parameter(torch.zeros(1, context_length, text_width))
-        layer = nn.TransformerEncoderLayer(
-            text_width,
-            text_heads,
-            text_width * 4,
-            dropout=0.0,
-            activation="gelu",
-            batch_first=True,
-            norm_first=True,
+        self.position_embedding = nn.Parameter(torch.empty(context_length, text_width))
+        text_layer = nn.TransformerEncoderLayer(
+            text_width, text_heads, text_width * 4, activation="gelu",
+            batch_first=True, norm_first=True, dropout=0.0,
         )
-        self.text_encoder = nn.TransformerEncoder(layer, text_layers)
-        self.text_projection = nn.Linear(text_width, embed_dim)
+        self.text_transformer = nn.TransformerEncoder(text_layer, text_layers)
+        self.text_norm = nn.LayerNorm(text_width)
+        self.text_projection = nn.Parameter(torch.empty(text_width, embed_dim))
         self.logit_scale = nn.Parameter(torch.tensor(math.log(1 / 0.07)))
         nn.init.normal_(self.position_embedding, std=0.01)
+        nn.init.normal_(self.text_projection, std=text_width ** -0.5)
 
     def encode_image(self, images):
-        return F.normalize(self.image_encoder(images), dim=-1)
+        if images.ndim != 4 or images.shape[1:] != (3, 224, 224):
+            raise ValueError("images must have paper-compatible shape [B,3,224,224]")
+        return F.normalize(self.visual(images), dim=-1)
 
     def encode_text(self, tokens):
-        features = self.text_encoder(self.token_embedding(tokens) + self.position_embedding)
-        mask = tokens.ne(0).unsqueeze(-1)
-        pooled = (features * mask).sum(1) / mask.sum(1).clamp_min(1)
-        return F.normalize(self.text_projection(pooled), dim=-1)
+        if tokens.ndim != 2 or tokens.shape[1] != self.context_length:
+            raise ValueError(f"tokens must have shape [B,{self.context_length}]")
+        causal_mask = torch.full(
+            (self.context_length, self.context_length), float("-inf"), device=tokens.device
+        ).triu_(1)
+        features = self.token_embedding(tokens) + self.position_embedding
+        features = self.text_norm(self.text_transformer(features, mask=causal_mask))
+        eot_positions = tokens.eq(self.eot_token_id).to(torch.int64).argmax(dim=-1)
+        pooled = features[torch.arange(tokens.shape[0], device=tokens.device), eot_positions]
+        return F.normalize(pooled @ self.text_projection, dim=-1)
 
     def forward(self, images, tokens):
-        image_features = self.encode_image(images)
-        text_features = self.encode_text(tokens)
-        scale = self.logit_scale.exp().clamp(max=100)
-        logits = scale * image_features @ text_features.t()
-        targets = torch.arange(logits.shape[0], device=logits.device)
-        loss = (
-            F.cross_entropy(logits, targets)
-            + F.cross_entropy(logits.t(), targets)
-        ) / 2
-        return {
-            "loss": loss,
-            "logits_per_image": logits,
-            "image_features": image_features,
-            "text_features": text_features,
-        }
+        return self.encode_image(images), self.encode_text(tokens), self.logit_scale.exp().clamp(max=100)
+
+
+def multi_positive_clip_loss(image_features, text_features, pair_ids, logit_scale):
+    """Symmetric CLIP loss where all samples sharing pair_id are positives."""
+    logits = logit_scale * image_features @ text_features.t()
+    positives = pair_ids[:, None].eq(pair_ids[None, :])
+    log_i = F.log_softmax(logits, dim=1)
+    log_t = F.log_softmax(logits.t(), dim=1)
+    loss_i = -(log_i.masked_fill(~positives, 0).sum(1) / positives.sum(1))
+    loss_t = -(log_t.masked_fill(~positives.t(), 0).sum(1) / positives.t().sum(1))
+    return (loss_i.mean() + loss_t.mean()) / 2
+
+
+__all__ = ["RemoteCLIP", "multi_positive_clip_loss"]

scripts/fake_data.py

@@ -1,54 +1,44 @@
-"""Generate deterministic RemoteCLIP-format image-text pairs."""
+"""Generate tiny paper-shape RemoteCLIP train/test NPZ files."""
 
+import json
 from pathlib import Path
 
 import numpy as np
 import yaml
 
-
 ROOT = Path(__file__).resolve().parents[1]
 
 
 def make_split(count, config, seed):
     rng = np.random.default_rng(seed)
-    data = config["data"]
-    size = data["image_size"]
-    classes = data["num_classes"]
-    images = np.empty((count, 3, size, size), dtype=np.float32)
-    tokens = np.zeros((count, data["context_length"]), dtype=np.int64)
-    labels = np.arange(count, dtype=np.int64) % classes
-    y, x = np.mgrid[0:size, 0:size].astype(np.float32) / max(size - 1, 1)
-    for index, label in enumerate(labels):
-        image = np.zeros((3, size, size), dtype=np.float32)
-        image[label % 3] = 0.55 + 0.35 * np.sin((label + 1) * np.pi * x)
-        image[(label + 1) % 3] += 0.25 * np.cos((label + 1) * np.pi * y)
-        images[index] = np.clip(image + rng.normal(0, 0.02, image.shape), 0, 1)
-        tokens[index, :4] = [label + 1, 16 + label, 32 + label, 48 + label]
-    return images, tokens, labels
+    d = config["data"]
+    images = np.empty((count, 3, 224, 224), dtype=np.float32)
+    tokens = np.zeros((count, 77), dtype=np.int64)
+    pair_ids = np.arange(count, dtype=np.int64) % d["num_semantic_groups"]
+    yy, xx = np.mgrid[:224, :224].astype(np.float32) / 223
+    for index, pair_id in enumerate(pair_ids):
+        base = np.stack((xx, yy, (xx + yy) / 2)) if pair_id == 0 else np.stack((yy, 1 - xx, xx * yy))
+        images[index] = np.clip(base + rng.normal(0, 0.025, base.shape), 0, 1)
+        # CLIP convention: EOT has the largest vocabulary id and therefore wins argmax pooling.
+        tokens[index, :6] = [49406, 100 + pair_id, 200 + pair_id, 300 + index, 400 + pair_id, 49407]
+    return images, tokens, pair_ids
 
 
 def main():
-    with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
-        config = yaml.safe_load(handle)
-    train = make_split(config["data"]["train_samples"], config, config["seed"])
-    test = make_split(config["data"]["test_samples"], config, config["seed"] + 1)
-    output = ROOT / config["data"]["path"]
-    output.parent.mkdir(parents=True, exist_ok=True)
-    np.savez_compressed(
-        output,
-        train_images=train[0],
-        train_tokens=train[1],
-        train_labels=train[2],
-        test_images=test[0],
-        test_tokens=test[1],
-        test_labels=test[2],
-        data_source=np.asarray("synthetic"),
-        protocol=np.asarray(config["data"]["protocol"]),
-    )
-    print(
-        f"generated={output.relative_to(ROOT)} train={len(train[0])} test={len(test[0])} "
-        f"data_source=synthetic protocol={config['data']['protocol']}"
-    )
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text())
+    output = ROOT / config["data"]["root"]
+    output.mkdir(parents=True, exist_ok=True)
+    for split, count, seed in (("train", config["data"]["train_samples"], config["seed"]),
+                               ("test", config["data"]["test_samples"], config["seed"] + 1)):
+        images, tokens, pair_ids = make_split(count, config, seed)
+        np.savez_compressed(output / f"{split}.npz", images=images, tokens=tokens, pair_ids=pair_ids,
+                            data_source=np.asarray("synthetic"), protocol=np.asarray(config["data"]["protocol"]))
+    (output / "format.json").write_text(json.dumps({
+        "protocol": config["data"]["protocol"], "data_source": "synthetic",
+        "images": "float32 [N,3,224,224] in [0,1]", "tokens": "int64 [N,77] CLIP BPE ids",
+        "pair_ids": "int64 [N], equal ids define valid multi-positive matches"
+    }, indent=2) + "\n")
+    print(f"created {output / 'train.npz'} and {output / 'test.npz'}")
 
 
 if __name__ == "__main__":

scripts/inference.py

@@ -1,66 +1,47 @@
-"""Extract RemoteCLIP features and image-text similarity scores."""
+"""Run checkpoint-backed RemoteCLIP retrieval inference."""
 
+import argparse
 import importlib.util
 from pathlib import Path
-
 import numpy as np
 import torch
 import yaml
 
-
 ROOT = Path(__file__).resolve().parents[1]
 
 
-def load_model_class():
-    spec = importlib.util.spec_from_file_location("remoteclip_model", ROOT / "model" / "remoteclip.py")
-    module = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(module)
-    return module.RemoteCLIP
-
-
 def main():
-    with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
-        config = yaml.safe_load(handle)
-    checkpoint_path = ROOT / config["paths"]["checkpoint"]
-    data_path = ROOT / config["data"]["path"]
-    if not checkpoint_path.exists():
-        raise FileNotFoundError("Missing checkpoint. Run `python scripts/train.py` first.")
-    if not data_path.exists():
-        raise FileNotFoundError("Missing data. Run `python scripts/fake_data.py` first.")
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
-    RemoteCLIP = load_model_class()
-    model = RemoteCLIP(
-        vocabulary_size=config["data"]["vocabulary_size"],
-        context_length=config["data"]["context_length"],
-        **config["model"],
-    ).to(device)
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
+    parser.add_argument("--data", type=Path); parser.add_argument("--checkpoint", type=Path)
+    parser.add_argument("--output-dir", type=Path); parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default="auto")
+    args = parser.parse_args(); config = yaml.safe_load(args.config.read_text())
+    checkpoint_path = args.checkpoint or ROOT / config["paths"]["checkpoint"]
+    if not checkpoint_path.is_file(): raise FileNotFoundError(f"checkpoint not found: {checkpoint_path}")
+    spec = importlib.util.spec_from_file_location("remoteclip", ROOT / "model/remoteclip.py")
+    module = importlib.util.module_from_spec(spec); spec.loader.exec_module(module)
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    model = module.RemoteCLIP(vocabulary_size=config["data"]["vocabulary_size"],
+                              context_length=config["data"]["context_length"],
+                              eot_token_id=config["data"]["eot_token_id"], **config["model"])
     model.load_state_dict(checkpoint["model"])
-    model.eval()
-    archive = np.load(data_path)
-    images = torch.from_numpy(archive["test_images"]).to(device)
-    tokens = torch.from_numpy(archive["test_tokens"]).to(device)
+    use_cuda = torch.cuda.is_available() and args.device != "cpu"
+    if args.device == "cuda" and not use_cuda: raise RuntimeError("CUDA requested but unavailable")
+    device = torch.device("cuda" if use_cuda else "cpu"); model.to(device).eval()
+    data_path = args.data or ROOT / config["data"]["root"] / "test.npz"
+    train_spec = importlib.util.spec_from_file_location("remoteclip_train", ROOT / "scripts/train.py")
+    train_module = importlib.util.module_from_spec(train_spec); train_spec.loader.exec_module(train_module)
+    dataset = train_module.PairDataset(data_path, config); archive = np.load(data_path)
     with torch.inference_mode():
-        image_features = model.encode_image(images)
-        text_features = model.encode_text(tokens)
-        similarities = image_features @ text_features.t()
-    output_dir = ROOT / config["paths"]["inference_dir"]
-    output_dir.mkdir(parents=True, exist_ok=True)
-    np.savez_compressed(
-        output_dir / "retrieval.npz",
-        similarities=similarities.cpu().numpy(),
-        image_features=image_features.cpu().numpy(),
-        text_features=text_features.cpu().numpy(),
-        labels=archive["test_labels"],
-        images=archive["test_images"],
-        data_source=archive["data_source"],
-        protocol=archive["protocol"],
-    )
-    print(
-        f"output={output_dir.relative_to(ROOT)} samples={len(images)} "
-        f"data_source={str(archive['data_source'])} protocol={str(archive['protocol'])}"
-    )
-
-
-if __name__ == "__main__":
-    main()
+        image_features = model.encode_image(torch.from_numpy(archive["images"]).to(device))
+        text_features = model.encode_text(torch.from_numpy(archive["tokens"]).to(device))
+    output_dir = args.output_dir or ROOT / config["paths"]["inference_dir"]; output_dir.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output_dir / "retrieval.npz", similarities=(image_features @ text_features.T).cpu().numpy(),
+                        image_features=image_features.cpu().numpy(), text_features=text_features.cpu().numpy(),
+                        pair_ids=archive["pair_ids"], images=archive["images"], checkpoint=np.asarray(str(checkpoint_path)),
+                        data_source=archive["data_source"] if "data_source" in archive else np.asarray("provided"),
+                        protocol=archive["protocol"] if "protocol" in archive else np.asarray("provided_npz"))
+    print(f"inference={output_dir / 'retrieval.npz'} checkpoint={checkpoint_path}")
+
+
+if __name__ == "__main__": main()

scripts/result.py

@@ -1,69 +1,41 @@
-"""Evaluate bidirectional retrieval and render a similarity matrix."""
+"""Evaluate multi-positive bidirectional retrieval and visualize similarities."""
 
+import argparse
 import json
 from pathlib import Path
-
 import matplotlib.pyplot as plt
 import numpy as np
 import yaml
 
-
 ROOT = Path(__file__).resolve().parents[1]
 
 
-def recall(similarities, k, transpose=False):
-    scores = similarities.T if transpose else similarities
-    topk = np.argsort(-scores, axis=1)[:, :k]
-    targets = np.arange(len(scores))[:, None]
-    return float((topk == targets).any(axis=1).mean())
+def recall(scores, query_ids, candidate_ids, k):
+    top = np.argsort(-scores, axis=1)[:, :min(k, scores.shape[1])]
+    return float(np.mean([np.isin(candidate_ids[index], query_ids[row]).any() for row, index in enumerate(top)]))
 
 
 def main():
-    with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
-        config = yaml.safe_load(handle)
-    input_path = ROOT / config["paths"]["inference_dir"] / "retrieval.npz"
-    if not input_path.exists():
-        raise FileNotFoundError("Missing inference output. Run `python scripts/inference.py` first.")
-    archive = np.load(input_path)
-    similarities = archive["similarities"]
-    limit = len(similarities)
-    metrics = {
-        "image_to_text_r1": recall(similarities, 1),
-        "image_to_text_r5": recall(similarities, min(5, limit)),
-        "text_to_image_r1": recall(similarities, 1, True),
-        "text_to_image_r5": recall(similarities, min(5, limit), True),
-        "mean_recall": 0.0,
-        "samples": limit,
-        "data_source": str(archive["data_source"]),
-        "protocol": str(archive["protocol"]),
-    }
-    metrics["mean_recall"] = float(
-        np.mean(
-            [
-                metrics["image_to_text_r1"],
-                metrics["image_to_text_r5"],
-                metrics["text_to_image_r1"],
-                metrics["text_to_image_r5"],
-            ]
-        )
-    )
-    output_dir = ROOT / config["paths"]["evaluation_dir"]
-    output_dir.mkdir(parents=True, exist_ok=True)
-    (output_dir / "metrics.json").write_text(
-        json.dumps(metrics, indent=2) + "\n", encoding="utf-8"
-    )
-    figure, axis = plt.subplots(figsize=(5, 4))
-    image = axis.imshow(similarities, cmap="viridis")
-    axis.set_xlabel("Text index")
-    axis.set_ylabel("Image index")
-    axis.set_title("RemoteCLIP image-text similarity")
-    figure.colorbar(image, ax=axis)
-    figure.tight_layout()
-    figure.savefig(output_dir / "similarity_matrix.png", dpi=120)
-    plt.close(figure)
-    print(json.dumps(metrics))
-    print(f"evaluation={output_dir.relative_to(ROOT)}")
-
-
-if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
+    parser.add_argument("--input", type=Path); parser.add_argument("--output-dir", type=Path)
+    args = parser.parse_args(); config = yaml.safe_load(args.config.read_text())
+    source = args.input or ROOT / config["paths"]["inference_dir"] / "retrieval.npz"
+    if not source.is_file(): raise FileNotFoundError("Run inference before evaluation")
+    archive = np.load(source); scores, ids = archive["similarities"], archive["pair_ids"]
+    metrics = {}
+    for k in (1, 5, 10):
+        metrics[f"image_to_text_R@{k}"] = recall(scores, ids, ids, k)
+        metrics[f"text_to_image_R@{k}"] = recall(scores.T, ids, ids, k)
+    metrics["mean_recall"] = float(np.mean(list(metrics.values())))
+    metrics.update(samples=int(len(ids)), protocol=str(archive["protocol"]), checkpoint=str(archive["checkpoint"]),
+                   multi_positive=True, data_source=str(archive["data_source"]))
+    output = args.output_dir or ROOT / config["paths"]["evaluation_dir"]; output.mkdir(parents=True, exist_ok=True)
+    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    figure, axis = plt.subplots(figsize=(5.4, 4.5)); image = axis.imshow(scores, cmap="magma")
+    axis.set(xlabel="Text candidate", ylabel="Image query", title="RemoteCLIP cosine similarity")
+    figure.colorbar(image, ax=axis); figure.tight_layout(); figure.savefig(output / "similarity_matrix.png", dpi=160); plt.close(figure)
+    print(json.dumps(metrics, indent=2)); print(f"evaluation={output}")
+
+
+if __name__ == "__main__": main()

scripts/train.py

@@ -1,132 +1,158 @@
-"""Train compact RemoteCLIP with bidirectional InfoNCE."""
+"""Train RemoteCLIP with multi-positive symmetric contrastive loss and torchrun."""
 
+import argparse
 import importlib.util
 import json
 import os
-import random
+from contextlib import nullcontext
+from functools import partial
 from pathlib import Path
 
 import numpy as np
 import torch
 import yaml
 from torch import distributed as dist
+from torch.distributed.nn import functional as dist_nn
 from torch.nn.parallel import DistributedDataParallel
 from torch.utils.data import DataLoader, Dataset, DistributedSampler
 
-
 ROOT = Path(__file__).resolve().parents[1]
 
 
-def load_model_class():
-    spec = importlib.util.spec_from_file_location("remoteclip_model", ROOT / "model" / "remoteclip.py")
+def load_model_module():
+    spec = importlib.util.spec_from_file_location("remoteclip", ROOT / "model/remoteclip.py")
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
-    return module.RemoteCLIP
+    return module
 
 
 class PairDataset(Dataset):
-    def __init__(self, path):
+    def __init__(self, path, config):
         archive = np.load(path)
-        self.images = archive["train_images"]
-        self.tokens = archive["train_tokens"]
-        self.data_source = str(archive["data_source"]) if "data_source" in archive.files else "unknown"
-        self.protocol = str(archive["protocol"]) if "protocol" in archive.files else "unknown"
+        required = {"images", "tokens", "pair_ids"}
+        if not required.issubset(archive.files):
+            raise ValueError(f"NPZ requires {sorted(required)}")
+        self.images, self.tokens, self.pair_ids = (archive[key] for key in ("images", "tokens", "pair_ids"))
+        self.data_source = str(archive["data_source"]) if "data_source" in archive else "provided"
+        self.protocol = str(archive["protocol"]) if "protocol" in archive else "provided_npz"
+        if self.protocol != config["data"]["protocol"]:
+            raise ValueError("NPZ protocol does not match configuration")
+        if self.images.ndim != 4 or self.images.shape[1:] != (3, 224, 224) or self.images.dtype != np.float32:
+            raise ValueError("images must be float32 [N,3,224,224]")
+        if self.tokens.ndim != 2 or self.tokens.shape[1:] != (77,) or self.tokens.dtype != np.int64:
+            raise ValueError("tokens must be int64 [N,77]")
+        if self.pair_ids.ndim != 1 or self.pair_ids.dtype != np.int64:
+            raise ValueError("pair_ids must be int64 [N]")
+        if not (len(self.images) == len(self.tokens) == len(self.pair_ids)) or len(self.images) == 0:
+            raise ValueError("images, tokens, and pair_ids must have the same non-zero sample count")
+        vocabulary_size = config["data"]["vocabulary_size"]
+        if self.tokens.min() < 0 or self.tokens.max() >= vocabulary_size:
+            raise ValueError(f"token ids must be in [0,{vocabulary_size})")
+        sot, eot, pad = (config["data"][key] for key in ("sot_token_id", "eot_token_id", "pad_token_id"))
+        if not np.all(self.tokens[:, 0] == sot):
+            raise ValueError("standard CLIP sequences must start with SOT")
+        eot_mask = self.tokens == eot
+        if not np.all(eot_mask.sum(axis=1) == 1):
+            raise ValueError("each token sequence must contain exactly one EOT")
+        eot_positions = eot_mask.argmax(axis=1)
+        for row, position in zip(self.tokens, eot_positions):
+            if np.any(row[1:position] == pad) or np.any(row[position + 1:] != pad):
+                raise ValueError("tokens before EOT must be non-padding and all tokens after EOT must be padding")
+        if not (0 <= pad < sot < eot < vocabulary_size):
+            raise ValueError("CLIP token configuration must satisfy pad < SOT < EOT < vocabulary_size")
 
-    def __len__(self):
-        return len(self.images)
+    def __len__(self): return len(self.images)
 
     def __getitem__(self, index):
-        return torch.from_numpy(self.images[index]), torch.from_numpy(self.tokens[index])
+        return tuple(torch.as_tensor(array[index]) for array in (self.images, self.tokens, self.pair_ids))
+
+
+def gather_with_local_grad(features):
+    if not dist.is_initialized():
+        return features
+    # Autograd all_gather uses reduce-scatter in backward. Combined with DDP's
+    # parameter-gradient averaging, this is the exact gradient of one global loss.
+    return torch.cat(dist_nn.all_gather(features), dim=0)
+
+
+def gather_ids(pair_ids):
+    if not dist.is_initialized():
+        return pair_ids
+    gathered = [torch.zeros_like(pair_ids) for _ in range(dist.get_world_size())]
+    dist.all_gather(gathered, pair_ids)
+    return torch.cat(gathered)
+
+
+def model_kwargs(config):
+    return {"vocabulary_size": config["data"]["vocabulary_size"],
+            "context_length": config["data"]["context_length"],
+            "eot_token_id": config["data"]["eot_token_id"], **config["model"]}
 
 
 def main():
-    with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
-        config = yaml.safe_load(handle)
-    data_path = ROOT / config["data"]["path"]
-    if not data_path.exists():
-        raise FileNotFoundError(
-            f"Missing training data: {data_path.relative_to(ROOT)}. "
-            "Run `python scripts/fake_data.py` first."
-        )
-    world_size = int(os.environ.get("WORLD_SIZE", "1"))
-    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml")
+    parser.add_argument("--data", type=Path)
+    parser.add_argument("--checkpoint", type=Path)
+    parser.add_argument("--device", choices=("auto", "cpu", "cuda"))
+    args = parser.parse_args()
+    config = yaml.safe_load(args.config.read_text())
+    world_size, rank = int(os.environ.get("WORLD_SIZE", 1)), int(os.environ.get("RANK", 0))
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    requested = args.device or config["runtime"]["device"]
+    use_cuda = torch.cuda.is_available() and requested != "cpu"
+    if requested == "cuda" and not use_cuda: raise RuntimeError("CUDA requested but unavailable")
+    if world_size > 1: dist.init_process_group("nccl" if use_cuda else "gloo")
+    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
+    if use_cuda: torch.cuda.set_device(local_rank)
+    torch.manual_seed(config["seed"] + rank)
+    data_path = args.data or ROOT / config["data"]["root"] / "train.npz"
+    dataset = PairDataset(data_path, config)
+    sampler = DistributedSampler(dataset, shuffle=True, drop_last=True) if world_size > 1 else None
+    loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], sampler=sampler,
+                        shuffle=sampler is None, num_workers=config["training"]["num_workers"],
+                        drop_last=world_size > 1)
+    module = load_model_module()
+    model = module.RemoteCLIP(**model_kwargs(config)).to(device)
+    raw_model = model
     if world_size > 1:
-        dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
-    device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() else "cpu")
-    if device.type == "cuda":
-        torch.cuda.set_device(local_rank)
-    seed = config["seed"] + local_rank
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    dataset = PairDataset(data_path)
-    sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None
-    loader = DataLoader(
-        dataset,
-        batch_size=config["train"]["batch_size"],
-        shuffle=sampler is None,
-        sampler=sampler,
-        num_workers=config["train"]["num_workers"],
-    )
-    RemoteCLIP = load_model_class()
-    model = RemoteCLIP(
-        vocabulary_size=config["data"]["vocabulary_size"],
-        context_length=config["data"]["context_length"],
-        **config["model"],
-    ).to(device)
-    if world_size > 1:
-        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
-    optimizer = torch.optim.AdamW(
-        model.parameters(),
-        lr=config["train"]["learning_rate"],
-        weight_decay=config["train"]["weight_decay"],
-    )
+        model = DistributedDataParallel(model, device_ids=[local_rank] if use_cuda else None)
+        raw_model = model.module
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config["training"]["learning_rate"],
+                                  weight_decay=config["training"]["weight_decay"])
+    amp = bool(config["runtime"]["amp"] and use_cuda)
+    scaler = torch.amp.GradScaler("cuda", enabled=amp)
     history = []
-    for epoch in range(config["train"]["epochs"]):
-        if sampler is not None:
-            sampler.set_epoch(epoch)
-        model.train()
-        total = 0.0
-        for images, tokens in loader:
-            output = model(images.to(device), tokens.to(device))
-            optimizer.zero_grad(set_to_none=True)
-            output["loss"].backward()
-            optimizer.step()
-            total += output["loss"].item()
-        loss = total / len(loader)
-        history.append({"epoch": epoch + 1, "contrastive_loss": loss})
-        if local_rank == 0:
-            print(f"epoch={epoch + 1} contrastive_loss={loss:.6f}")
-    if local_rank == 0:
-        checkpoint = ROOT / config["paths"]["checkpoint"]
-        metrics = ROOT / config["paths"]["training_metrics"]
-        checkpoint.parent.mkdir(parents=True, exist_ok=True)
-        metrics.parent.mkdir(parents=True, exist_ok=True)
-        base_model = model.module if hasattr(model, "module") else model
-        torch.save(
-            {
-                "model": base_model.state_dict(),
-                "config": config,
-                "data_source": dataset.data_source,
-                "protocol": dataset.protocol,
-            },
-            checkpoint,
-        )
-        metrics.write_text(
-            json.dumps(
-                {"history": history, "data_source": dataset.data_source, "protocol": dataset.protocol},
-                indent=2,
-            )
-            + "\n",
-            encoding="utf-8",
-        )
-        print(
-            f"checkpoint={checkpoint.relative_to(ROOT)} data_source={dataset.data_source} "
-            f"protocol={dataset.protocol}"
-        )
-    if world_size > 1:
-        dist.destroy_process_group()
+    for epoch in range(config["training"]["epochs"]):
+        if sampler is not None: sampler.set_epoch(epoch)
+        model.train(); total = 0.0
+        for images, tokens, pair_ids in loader:
+            images, tokens, pair_ids = images.to(device), tokens.to(device), pair_ids.to(device)
+            autocast = partial(torch.amp.autocast, "cuda") if amp else nullcontext
+            with autocast():
+                image_features, text_features, scale = model(images, tokens)
+                loss = module.multi_positive_clip_loss(gather_with_local_grad(image_features),
+                                                       gather_with_local_grad(text_features),
+                                                       gather_ids(pair_ids), scale)
+            optimizer.zero_grad(set_to_none=True); scaler.scale(loss).backward()
+            scaler.step(optimizer); scaler.update(); total += loss.detach().item()
+        statistics = torch.tensor([total, len(loader)], dtype=torch.float64, device=device)
+        if world_size > 1: dist.all_reduce(statistics, op=dist.ReduceOp.SUM)
+        record = {"epoch": epoch + 1, "contrastive_loss": statistics[0].item() / max(statistics[1].item(), 1)}
+        history.append(record)
+        if rank == 0: print(f"epoch={record['epoch']} contrastive_loss={record['contrastive_loss']:.6f}")
+    if rank == 0:
+        checkpoint_path = args.checkpoint or ROOT / config["paths"]["checkpoint"]
+        checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({"model": raw_model.state_dict(), "optimizer": optimizer.state_dict(),
+                    "scaler": scaler.state_dict() if amp else None, "config": config,
+                    "epoch": config["training"]["epochs"], "history": history}, checkpoint_path)
+        metrics = ROOT / config["paths"]["training_metrics"]; metrics.parent.mkdir(parents=True, exist_ok=True)
+        metrics.write_text(json.dumps({"history": history, "protocol": dataset.protocol,
+                                       "data_source": dataset.data_source,
+                                       "world_size": world_size, "amp": amp}, indent=2) + "\n")
+        print(f"checkpoint={checkpoint_path}")
+    if world_size > 1: dist.destroy_process_group()
 
 
-if __name__ == "__main__":
-    main()
+if __name__ == "__main__": main()