Add RemoteCLIP paper reproduction package

.gitattributes

@@ -1,35 +1,13 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
 *.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,183 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Remote Sensing
+- Vision-Language Model
+- Image-Text Retrieval
+- RemoteCLIP
+- arxiv:2306.11029
+tasks: []
+datasets:
+- RSITMD
+- RSICD
+- UCM-Captions
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">RemoteCLIP</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+RemoteCLIP is a vision-language foundation model for remote sensing. It learns semantically rich visual representations by aligning remote sensing images with text descriptions and supports image-text retrieval, zero-shot classification, few-shot classification, linear probing, k-NN classification, and remote sensing object counting.
+
+Paper: RemoteCLIP: A Vision Language Foundation Model for Remote Sensing
+
+https://arxiv.org/abs/2306.11029
+
+# Model Description
+
+RemoteCLIP was proposed by researchers from Hohai University, the Hong Kong University of Science and Technology, Nanjing Forestry University, the Chinese Academy of Forestry, and Griffith University. It is continually pretrained on RSITMD, RSICD, UCM-Captions, and remote sensing detection, segmentation, and UAV imagery converted into image-text pairs through B2C and M2B, and is suitable for remote sensing image-text retrieval and open-vocabulary recognition.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Remote sensing image-text contrastive training | Train RemoteCLIP with paired remote sensing images and text descriptions. |
+| Remote sensing image-text retrieval | Retrieve text from images or remote sensing images from text queries. |
+| Local quick validation | Use synthetic image-text pairs to validate training, inference, evaluation, and visualization. |
+| Hugging Face / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Launch distributed multi-process training with `torchrun`. |
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching the cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/RemoteCLIP --local-dir ./RemoteCLIP
+cd RemoteCLIP
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Please activate CONDA first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data Introduction
+
+The paper unifies three groups of heterogeneous remote sensing datasets into paired image-text data. RET-3 contains RSITMD, RSICD, and UCM-Captions; DET-10 contains satellite and UAV object detection datasets; SEG-4 contains the iSAID, LoveDA, Potsdam, and Vaihingen semantic segmentation datasets. Detection annotations are converted into descriptions through Box-to-Caption, while segmentation masks are converted to boxes through Mask-to-Box before caption generation. The resulting pretraining collection is approximately 12 times larger than the original remote sensing image-text datasets combined.
+
+The paper datasets remain subject to their respective download and usage licenses. The authors' official code and model resources are available at:
+
+```text
+https://github.com/ChenDelong1999/RemoteCLIP
+```
+
+The complete paper pretraining dataset is not included in this model repository. Generate compact synthetic data for local validation with:
+
+```bash
+python scripts/fake_data.py
+```
+
+This command creates 32 training pairs and 8 test pairs in `data/remoteclip.npz`. Synthetic data only validates image and text encoding, bidirectional InfoNCE training, checkpoint loading, retrieval inference, evaluation, and similarity-matrix visualization. The complete synthetic pipeline has been verified successfully, but synthetic training values are not reported as model performance.
+
+For real-data training, organize remote sensing images and text descriptions into paired samples following the paper or the official repository and convert them to `data/remoteclip.npz`. Do not run `scripts/fake_data.py` for real-data training. Set `data.protocol` in `conf/config.yaml` to `real_remoteclip` and update the image size, vocabulary size, context length, and training parameters for the converted dataset.
+
+### Training
+
+Generate synthetic data before local quick validation:
+
+```bash
+python scripts/fake_data.py
+```
+
+Single GPU or CPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+Training outputs:
+
+```text
+result/checkpoints/remoteclip.pt
+result/training/metrics.json
+```
+
+### Training Weights
+
+Pretrained weights based on paired remote sensing image-text data will be provided under `weight/`. The weight file will be uploaded in a future update. Checkpoints produced from synthetic data validate the pipeline only and do not provide real remote sensing image-text retrieval capability.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference extracts image features, text features, and the cross-modal similarity matrix from the test pairs.
+
+Prediction output:
+
+```text
+result/output/retrieval.npz
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation outputs:
+
+```text
+result/evaluation/metrics.json
+result/evaluation/similarity_matrix.png
+```
+
+The evaluation script computes image-to-text and text-to-image R@1, R@5, and mean recall, and renders the image-text similarity matrix. Synthetic data is used only to confirm that evaluation and visualization outputs can be generated successfully; its numerical values are not reported as model performance.
+
+# Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- This repository is a reproduction of the original RemoteCLIP paper.
+- The reproduction code is released under Apache License 2.0. The original datasets and official model weights remain subject to their respective source licenses.

conf/config.yaml

@@ -0,0 +1,27 @@
+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
+model:
+  embed_dim: 32
+  vision_width: 32
+  text_width: 32
+  text_layers: 2
+  text_heads: 4
+train:
+  epochs: 5
+  batch_size: 8
+  learning_rate: 0.001
+  weight_decay: 0.0001
+  num_workers: 0
+paths:
+  checkpoint: result/checkpoints/remoteclip.pt
+  training_metrics: result/training/metrics.json
+  inference_dir: result/output
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,50 @@
+{
+  "model_name": "RemoteCLIP",
+  "model_type": "remoteclip",
+  "architectures": [
+    "RemoteCLIP"
+  ],
+  "framework": "PyTorch",
+  "domain": "earth-science",
+  "task": "remote-sensing-image-text-retrieval",
+  "implementation": {
+    "entry_point": "model/remoteclip.py",
+    "scope": "compact image-text contrastive RemoteCLIP reproduction"
+  },
+  "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,
+    "text_layers": 2,
+    "text_heads": 4,
+    "training_objective": "bidirectional InfoNCE"
+  },
+  "data": {
+    "datasets": [
+      "RSITMD",
+      "RSICD",
+      "UCM-Captions"
+    ],
+    "official_repo": "ChenDelong1999/RemoteCLIP",
+    "protocol": "synthetic_remoteclip",
+    "synthetic_train_samples": 32,
+    "synthetic_test_samples": 8
+  },
+  "metrics": [
+    "image_to_text_r1",
+    "image_to_text_r5",
+    "text_to_image_r1",
+    "text_to_image_r5",
+    "mean_recall"
+  ],
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/remoteclip.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "remote_sensing_image_text_retrieval",
+  "model": "RemoteCLIP",
+  "input_format": "image_text_npz",
+  "protocol": "synthetic_remoteclip",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/remoteclip.py

@@ -0,0 +1,72 @@
+"""Compact RemoteCLIP image-text contrastive model."""
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class RemoteCLIP(nn.Module):
+    def __init__(
+        self,
+        vocabulary_size=64,
+        context_length=8,
+        embed_dim=32,
+        vision_width=32,
+        text_width=32,
+        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.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.text_encoder = nn.TransformerEncoder(layer, text_layers)
+        self.text_projection = nn.Linear(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)
+
+    def encode_image(self, images):
+        return F.normalize(self.image_encoder(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)
+
+    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,
+        }

scripts/fake_data.py

@@ -0,0 +1,55 @@
+"""Generate deterministic RemoteCLIP-format image-text pairs."""
+
+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
+
+
+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']}"
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,66 @@
+"""Extract RemoteCLIP features and image-text similarity scores."""
+
+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)
+    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)
+    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()

scripts/result.py

@@ -0,0 +1,69 @@
+"""Evaluate bidirectional retrieval and render a similarity matrix."""
+
+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 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()

scripts/train.py

@@ -0,0 +1,132 @@
+"""Train compact RemoteCLIP with bidirectional InfoNCE."""
+
+import importlib.util
+import json
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch import distributed as dist
+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")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.RemoteCLIP
+
+
+class PairDataset(Dataset):
+    def __init__(self, path):
+        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"
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, index):
+        return torch.from_numpy(self.images[index]), torch.from_numpy(self.tokens[index])
+
+
+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"))
+    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"],
+    )
+    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()
+
+
+if __name__ == "__main__":
+    main()