Add resumable SigLIP2 Giant formal training

.gitea/workflows/gpu-train.yaml

@@ -0,0 +1,73 @@
+name: Formal GPU training
+
+on:
+  workflow_dispatch:
+
+jobs:
+  extractor:
+    runs-on: gpu
+    timeout-minutes: 1440
+    env:
+      HF_ENDPOINT: https://hf-mirror.com
+      HF_HUB_DISABLE_XET: "1"
+      TOKENIZERS_PARALLELISM: "false"
+    steps:
+      - name: Checkout from Gitea
+        run: |
+          git init .
+          git remote add origin "$GITHUB_SERVER_URL/$GITHUB_REPOSITORY.git"
+          git fetch --depth 1 origin "$GITHUB_SHA"
+          git checkout --detach FETCH_HEAD
+
+      - name: Experiment details
+        run: |
+          echo "::group::Hardware"
+          nvidia-smi
+          python -c "import torch; print(torch.__version__, torch.version.cuda, torch.cuda.get_device_name(), torch.cuda.get_device_capability())"
+          echo "::endgroup::"
+          echo "::group::Configuration"
+          cat configs/siglip2_giant.yaml
+          echo "::endgroup::"
+
+      - name: Prepare formal environment
+        run: |
+          source /etc/network_turbo
+          python -m venv --system-site-packages /root/autodl-tmp/ces-venv
+          /root/autodl-tmp/ces-venv/bin/python -m pip install numpy pillow pyyaml "transformers>=4.48,<5" accelerate
+          unset http_proxy https_proxy
+
+      - name: Build full synthetic dataset
+        run: |
+          /root/autodl-tmp/ces-venv/bin/python -m data.build_splits \
+            --config configs/siglip2_giant.yaml \
+            --reuse-existing
+
+      - name: Train SigLIP2 Giant extractor
+        run: |
+          source /etc/network_turbo
+          /root/autodl-tmp/ces-venv/bin/python train_extractor.py \
+            --config configs/siglip2_giant.yaml \
+            --backbone siglip2
+          unset http_proxy https_proxy
+
+      - name: Enforce research gates
+        run: |
+          /root/autodl-tmp/ces-venv/bin/python -m eval.benchmark_extractor \
+            --config configs/siglip2_giant.yaml \
+            --checkpoint /root/autodl-tmp/ces-artifacts/siglip2-giant/best.pt \
+            --backbone siglip2 \
+            --output /root/autodl-tmp/ces-artifacts/siglip2-giant/benchmark.json \
+            --enforce-gates
+
+      - name: Details
+        if: always()
+        run: |
+          echo "::group::Epoch details"
+          test -f /root/autodl-tmp/ces-artifacts/siglip2-giant/training_details.jsonl && cat /root/autodl-tmp/ces-artifacts/siglip2-giant/training_details.jsonl || true
+          echo "::endgroup::"
+          echo "::group::Benchmark details"
+          test -f /root/autodl-tmp/ces-artifacts/siglip2-giant/benchmark.json && cat /root/autodl-tmp/ces-artifacts/siglip2-giant/benchmark.json || true
+          echo "::endgroup::"
+          echo "::group::Saved artifacts"
+          ls -lh /root/autodl-tmp/ces-artifacts/siglip2-giant || true
+          echo "::endgroup::"

README.md

@@ -62,6 +62,35 @@ uv run python -m data.build_splits --config configs/mvp.yaml
 uv run python train_extractor.py --config configs/mvp.yaml --backbone clip
 ```
 
+### 正式 2B 级训练
+
+正式实验使用 `google/siglip2-giant-opt-patch16-384`：约 2B 总参数、384 分辨率的
+原生图文双塔。backbone 使用 BF16 并冻结，四层 512 维 Evidence Extractor 参与训练。
+`configs/siglip2_giant.yaml` 定义 30,000/3,000/3,000/3,000 数据划分与 10 epochs。
+
+Gitea 的 `Formal GPU training` workflow 会：
+
+- 使用 AutoDL `network_turbo` 与 `hf-mirror.com` 下载模型；
+- 在 `/root/autodl-tmp/ces-data/synthetic-v1` 复用完整数据集；
+- 每个 epoch 保存可恢复的 `last.pt`，并按验证集保存 `best.pt`；
+- 将运行配置、硬件、分项 loss、每轮耗时、所有验证指标和最终 benchmark 放在
+  Actions 的 `Details` step；
+- 将结果持久化到 `/root/autodl-tmp/ces-artifacts/siglip2-giant`；
+- 最后用六项研究门槛决定 workflow 是否通过。
+
+手动等价命令为：
+
+```bash
+python -m data.build_splits --config configs/siglip2_giant.yaml --reuse-existing
+python train_extractor.py --config configs/siglip2_giant.yaml --backbone siglip2
+python -m eval.benchmark_extractor \
+  --config configs/siglip2_giant.yaml \
+  --checkpoint /root/autodl-tmp/ces-artifacts/siglip2-giant/best.pt \
+  --backbone siglip2 \
+  --output /root/autodl-tmp/ces-artifacts/siglip2-giant/benchmark.json \
+  --enforce-gates
+```
+
 ## 每次训练后的独立验收
 
 以下命令运行统一 benchmark，写出 JSON，并在任一指标未达标时返回非零退出码：
@@ -82,7 +111,7 @@ uv run python -m eval.benchmark_extractor \
 | Cross-modal Recall@1 | 0.80 | 文本 FACT 检索全部视觉 FACT，最近邻标签完全相同 |
 | Hard-negative Recall@1 | 0.75 | 候选仅含正确项与恰好错一个基础字段的视觉 FACT |
 | OOD fact F1 | 0.80 | held-out 颜色–形状组合上的视觉 FACT exact-match F1 |
-| NULL slot accuracy | 0.95 | permutation matching 后，全部视觉/文本 slots 的类型准确率，含 NULL |
+| NULL slot accuracy | 0.95 | permutation matching 后，应为 NULL 的空槽中预测为 NULL 的比例 |
 
 CI 的单 epoch smoke run 会执行 benchmark 来验证整条评测管线，但不会使用
 `--enforce-gates`，因为 smoke 模型并不代表正式实验是否达标。正式训练完成后使用上述强制命令。

configs/mvp.yaml

@@ -12,6 +12,7 @@ data:
 model:
   backbone: tiny
   clip_model: openai/clip-vit-base-patch32
+  siglip2_model: google/siglip2-giant-opt-patch16-384
   state_dim: 256
   num_visual_slots: 6
   num_text_slots: 3

configs/prototype.yaml

@@ -12,6 +12,7 @@ data:
 model:
   backbone: tiny
   clip_model: openai/clip-vit-base-patch32
+  siglip2_model: google/siglip2-giant-opt-patch16-384
   state_dim: 256
   num_visual_slots: 6
   num_text_slots: 3

configs/siglip2_giant.yaml

@@ -0,0 +1,46 @@
+seed: 17
+
+data:
+  root: /root/autodl-tmp/ces-data/synthetic-v1
+  image_size: 384
+  train_size: 30000
+  validation_size: 3000
+  test_id_size: 3000
+  test_ood_size: 3000
+  num_objects: 2
+
+model:
+  backbone: siglip2
+  clip_model: openai/clip-vit-base-patch32
+  siglip2_model: google/siglip2-giant-opt-patch16-384
+  state_dim: 512
+  num_visual_slots: 6
+  num_text_slots: 3
+  extractor_layers: 4
+  num_heads: 8
+  dropout: 0.1
+
+training:
+  batch_size: 16
+  learning_rate: 0.0001
+  weight_decay: 0.01
+  epochs: 10
+  num_workers: 4
+  type_weight: 0.5
+  alignment_weight: 0.2
+  contrastive_temperature: 0.1
+  output_dir: /root/autodl-tmp/ces-artifacts/siglip2-giant
+  auto_resume: true
+
+prototype:
+  enabled: false
+  weight: 0.2
+
+evaluation:
+  gates:
+    text_fact_rule_f1: 0.98
+    visual_fact_f1: 0.90
+    cross_modal_recall_at_1: 0.80
+    hard_negative_recall_at_1: 0.75
+    ood_fact_f1: 0.80
+    null_slot_accuracy: 0.95

data/build_splits.py

@@ -33,7 +33,7 @@ def _write_split(
             handle.write(json.dumps(record, ensure_ascii=False) + "\n")
 
 
-def build(config_path: Path, small: bool = False) -> Path:
+def build(config_path: Path, small: bool = False, reuse_existing: bool = False) -> Path:
     config = yaml.safe_load(config_path.read_text(encoding="utf-8"))
     root = Path(config["data"]["root"])
     root.mkdir(parents=True, exist_ok=True)
@@ -45,6 +45,12 @@ def build(config_path: Path, small: bool = False) -> Path:
     }
     if small:
         sizes = {"train": 64, "validation": 16, "test_id": 16, "test_ood": 16}
+    if reuse_existing and all(
+        (root / f"{split}.jsonl").exists()
+        and sum(1 for _ in (root / f"{split}.jsonl").open(encoding="utf-8")) == size
+        for split, size in sizes.items()
+    ):
+        return root
     cursor = 0
     for split, size in sizes.items():
         _write_split(
@@ -64,8 +70,9 @@ def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml"))
     parser.add_argument("--small", action="store_true")
+    parser.add_argument("--reuse-existing", action="store_true")
     args = parser.parse_args()
-    print(build(args.config, small=args.small))
+    print(build(args.config, small=args.small, reuse_existing=args.reuse_existing))
 
 
 if __name__ == "__main__":

eval/benchmark_extractor.py

@@ -140,7 +140,7 @@ def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml"))
     parser.add_argument("--checkpoint", type=Path, required=True)
-    parser.add_argument("--backbone", choices=("tiny", "clip"), default=None)
+    parser.add_argument("--backbone", choices=("tiny", "clip", "siglip2"), default=None)
     parser.add_argument("--output", type=Path, default=None)
     parser.add_argument("--enforce-gates", action="store_true")
     args = parser.parse_args()
@@ -148,7 +148,7 @@ def main() -> None:
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = EvidenceExtractorModel(config, backbone_name=args.backbone).to(device)
     checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=True)
-    model.load_state_dict(checkpoint["model"])
+    model.load_state_dict(checkpoint["model"], strict=False)
     result = benchmark(model, config, device)
     rendered = json.dumps(result, indent=2)
     print(rendered)

eval/evaluate_slots.py

@@ -90,19 +90,24 @@ def summarize(counts: dict[str, int]) -> dict[str, float]:
     }
 
 
-def null_slot_counts(
+def slot_type_counts(
     predictions: dict[str, torch.Tensor], targets: list[torch.Tensor]
-) -> tuple[int, int]:
+) -> tuple[int, int, int, int]:
     matching = hungarian_matching_loss(predictions, targets)
     predicted_types = predictions["type_logits"].argmax(-1)
-    correct = 0
-    total = predicted_types.numel()
+    type_correct = 0
+    type_total = predicted_types.numel()
+    null_correct = 0
+    null_total = 0
     for batch_index, (slots, target_indices) in enumerate(matching.assignments):
         expected = torch.full_like(predicted_types[batch_index], NULL_TYPE)
         if len(slots):
             expected[slots] = targets[batch_index].to(expected.device)[target_indices, 0]
-        correct += int((predicted_types[batch_index] == expected).sum())
-    return correct, total
+        type_correct += int((predicted_types[batch_index] == expected).sum())
+        null_positions = expected == NULL_TYPE
+        null_correct += int((predicted_types[batch_index][null_positions] == NULL_TYPE).sum())
+        null_total += int(null_positions.sum())
+    return null_correct, null_total, type_correct, type_total
 
 
 @torch.no_grad()
@@ -118,6 +123,8 @@ def evaluate(
     }
     null_correct = 0
     null_total = 0
+    type_correct = 0
+    type_total = 0
     for batch in loader:
         outputs = model(batch["images"].to(device), batch["rules_text"])
         requests = (
@@ -134,11 +141,14 @@ def evaluate(
             ("visual", "visual_targets"),
             ("text", "text_targets"),
         ):
-            correct, count = null_slot_counts(outputs[modality], batch[target_name])
-            null_correct += correct
-            null_total += count
+            counts = slot_type_counts(outputs[modality], batch[target_name])
+            null_correct += counts[0]
+            null_total += counts[1]
+            type_correct += counts[2]
+            type_total += counts[3]
     result = {name: summarize(counts) for name, counts in totals.items()}
     result["null_slot_accuracy"] = null_correct / max(1, null_total)
+    result["slot_type_accuracy"] = type_correct / max(1, type_total)
     return result
 
 
@@ -146,14 +156,14 @@ def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml"))
     parser.add_argument("--checkpoint", type=Path, required=True)
-    parser.add_argument("--backbone", choices=("tiny", "clip"), default=None)
+    parser.add_argument("--backbone", choices=("tiny", "clip", "siglip2"), default=None)
     parser.add_argument("--split", default="test_id")
     args = parser.parse_args()
     config = yaml.safe_load(args.config.read_text(encoding="utf-8"))
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = EvidenceExtractorModel(config, backbone_name=args.backbone).to(device)
     state = torch.load(args.checkpoint, map_location=device, weights_only=True)
-    model.load_state_dict(state["model"])
+    model.load_state_dict(state["model"], strict=False)
     dataset = EvidenceDataset(
         Path(config["data"]["root"]) / f"{args.split}.jsonl",
         config["data"]["image_size"],

models/evidence_model.py

@@ -17,7 +17,12 @@ class EvidenceExtractorModel(nn.Module):
         model_config = config["model"]
         selected = backbone_name or model_config["backbone"]
         self.backbone_name = selected
-        self.backbone = build_backbone(selected, model_config["clip_model"])
+        model_name = (
+            model_config["siglip2_model"]
+            if selected == "siglip2"
+            else model_config["clip_model"]
+        )
+        self.backbone = build_backbone(selected, model_name)
         common = {
             "state_dim": model_config["state_dim"],
             "num_layers": model_config["extractor_layers"],

models/frozen_backbone.py

@@ -102,9 +102,71 @@ class CLIPFrozenBackbone(FrozenBackbone):
         return tokens, values["attention_mask"].bool()
 
 
-def build_backbone(name: str, clip_model: str) -> FrozenBackbone:
+class SigLIP2FrozenBackbone(FrozenBackbone):
+    """Billion-scale paired image/text encoder used for formal experiments."""
+
+    def __init__(self, model_name: str):
+        super().__init__()
+        try:
+            from transformers import AutoModel, AutoProcessor
+        except ImportError as error:
+            raise RuntimeError("Install the 'hf' extra to use SigLIP2.") from error
+        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
+        self.model = AutoModel.from_pretrained(
+            model_name,
+            torch_dtype=dtype,
+            low_cpu_mem_usage=True,
+        )
+        self.processor = AutoProcessor.from_pretrained(model_name)
+        self.model.requires_grad_(False).eval()
+        self.vision_dim = self.model.config.vision_config.hidden_size
+        self.text_dim = self.model.config.text_config.hidden_size
+
+    def _to_model_device(self, values: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        device = next(self.model.parameters()).device
+        dtype = next(self.model.parameters()).dtype
+        moved = {}
+        for key, value in values.items():
+            moved[key] = (
+                value.to(device=device, dtype=dtype)
+                if value.is_floating_point()
+                else value.to(device=device)
+            )
+        return moved
+
+    def encode_image(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        values = self.processor(images=images.detach().cpu(), return_tensors="pt")
+        allowed = {"pixel_values", "pixel_attention_mask", "spatial_shapes"}
+        vision_values = self._to_model_device(
+            {key: value for key, value in values.items() if key in allowed}
+        )
+        with torch.no_grad():
+            tokens = self.model.vision_model(**vision_values).last_hidden_state
+        mask = torch.ones(tokens.shape[:2], dtype=torch.bool, device=tokens.device)
+        return tokens, mask
+
+    def encode_text(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
+        max_length = self.model.config.text_config.max_position_embeddings
+        values = self.processor(
+            text=texts,
+            padding="max_length",
+            truncation=True,
+            max_length=max_length,
+            return_tensors="pt",
+        )
+        text_values = self._to_model_device(
+            {key: value for key, value in values.items() if key in {"input_ids", "attention_mask"}}
+        )
+        with torch.no_grad():
+            tokens = self.model.text_model(**text_values).last_hidden_state
+        return tokens, text_values["attention_mask"].bool()
+
+
+def build_backbone(name: str, model_name: str) -> FrozenBackbone:
     if name == "tiny":
         return TinyFrozenBackbone()
     if name == "clip":
-        return CLIPFrozenBackbone(clip_model)
+        return CLIPFrozenBackbone(model_name)
+    if name == "siglip2":
+        return SigLIP2FrozenBackbone(model_name)
     raise ValueError(f"Unknown backbone: {name}")

train_extractor.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import argparse
 import json
 import random
+import time
 from pathlib import Path
 from typing import Any
 
@@ -27,7 +28,17 @@ def set_seed(seed: int) -> None:
         torch.cuda.manual_seed_all(seed)
 
 
-def train(config: dict[str, Any], backbone_name: str, epochs: int | None = None) -> Path:
+def trainable_state_dict(model: EvidenceExtractorModel) -> dict[str, torch.Tensor]:
+    names = {name for name, parameter in model.named_parameters() if parameter.requires_grad}
+    return {name: value for name, value in model.state_dict().items() if name in names}
+
+
+def train(
+    config: dict[str, Any],
+    backbone_name: str,
+    epochs: int | None = None,
+    resume: Path | None = None,
+) -> Path:
     set_seed(config["seed"])
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     data_root = Path(config["data"]["root"])
@@ -54,12 +65,51 @@ def train(config: dict[str, Any], backbone_name: str, epochs: int | None = None)
         output_dir = output_dir.parent / backbone_name
     output_dir.mkdir(parents=True, exist_ok=True)
     best_path = output_dir / "best.pt"
+    last_path = output_dir / "last.pt"
+    details_path = output_dir / "training_details.jsonl"
     best_score = -1.0
     epoch_count = epochs or config["training"]["epochs"]
+    start_epoch = 1
+    resume_path = resume
+    if resume_path is None and config["training"].get("auto_resume") and last_path.exists():
+        resume_path = last_path
+    if resume_path is not None:
+        checkpoint = torch.load(resume_path, map_location=device, weights_only=True)
+        model.load_state_dict(checkpoint["model"], strict=False)
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        start_epoch = int(checkpoint["epoch"]) + 1
+        best_score = float(checkpoint.get("best_score", -1.0))
 
-    for epoch in range(1, epoch_count + 1):
+    total_parameters = sum(parameter.numel() for parameter in model.parameters())
+    trainable_parameters = sum(parameter.numel() for parameter in trainable)
+    run_details = {
+        "event": "run_start",
+        "device": str(device),
+        "gpu": torch.cuda.get_device_name() if torch.cuda.is_available() else None,
+        "backbone": backbone_name,
+        "total_parameters": total_parameters,
+        "trainable_parameters": trainable_parameters,
+        "train_samples": len(train_dataset),
+        "validation_samples": len(validation_dataset),
+        "start_epoch": start_epoch,
+        "target_epochs": epoch_count,
+        "config": config,
+    }
+    details_mode = "a" if start_epoch > 1 else "w"
+    with details_path.open(details_mode, encoding="utf-8") as handle:
+        handle.write(json.dumps(run_details) + "\n")
+    print(json.dumps(run_details))
+
+    for epoch in range(start_epoch, epoch_count + 1):
+        epoch_started = time.perf_counter()
         model.train()
-        running_loss = 0.0
+        running = {
+            "total": 0.0,
+            "visual": 0.0,
+            "text": 0.0,
+            "alignment": 0.0,
+            "prototype": 0.0,
+        }
         for batch in train_loader:
             optimizer.zero_grad(set_to_none=True)
             outputs = model(batch["images"].to(device), batch["rules_text"])
@@ -79,6 +129,7 @@ def train(config: dict[str, Any], backbone_name: str, epochs: int | None = None)
                 config["training"]["contrastive_temperature"],
             )
             loss = visual.loss + text.loss + config["training"]["alignment_weight"] * alignment
+            prototype_loss = loss.new_zeros(())
             if model.prototype_encoder is not None:
                 prototype_loss = canonical_prototype_loss(
                     outputs["visual"]["states"],
@@ -95,32 +146,61 @@ def train(config: dict[str, Any], backbone_name: str, epochs: int | None = None)
             loss.backward()
             torch.nn.utils.clip_grad_norm_(trainable, 1.0)
             optimizer.step()
-            running_loss += float(loss.detach())
+            running["total"] += float(loss.detach())
+            running["visual"] += float(visual.loss.detach())
+            running["text"] += float(text.loss.detach())
+            running["alignment"] += float(alignment.detach())
+            running["prototype"] += float(prototype_loss.detach())
         metrics = evaluate(model, validation_loader, device)
         score = metrics["visual_fact"]["f1"] + metrics["text_all"]["f1"]
         summary = {
+            "event": "epoch_end",
             "epoch": epoch,
-            "loss": running_loss / max(1, len(train_loader)),
+            "duration_seconds": time.perf_counter() - epoch_started,
+            "learning_rate": optimizer.param_groups[0]["lr"],
+            "loss": {
+                key: value / max(1, len(train_loader)) for key, value in running.items()
+            },
             "metrics": metrics,
         }
         print(json.dumps(summary))
+        with details_path.open("a", encoding="utf-8") as handle:
+            handle.write(json.dumps(summary) + "\n")
         if score > best_score:
             best_score = score
             torch.save(
-                {"model": model.state_dict(), "epoch": epoch, "metrics": metrics}, best_path
+                {
+                    "model": trainable_state_dict(model),
+                    "epoch": epoch,
+                    "metrics": metrics,
+                    "backbone": backbone_name,
+                },
+                best_path,
             )
+        torch.save(
+            {
+                "model": trainable_state_dict(model),
+                "optimizer": optimizer.state_dict(),
+                "epoch": epoch,
+                "metrics": metrics,
+                "best_score": best_score,
+                "backbone": backbone_name,
+            },
+            last_path,
+        )
     return best_path
 
 
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml"))
-    parser.add_argument("--backbone", choices=("tiny", "clip"), default=None)
+    parser.add_argument("--backbone", choices=("tiny", "clip", "siglip2"), default=None)
     parser.add_argument("--epochs", type=int, default=None)
+    parser.add_argument("--resume", type=Path, default=None)
     args = parser.parse_args()
     config = yaml.safe_load(args.config.read_text(encoding="utf-8"))
     backbone_name = args.backbone or config["model"]["backbone"]
-    print(train(config, backbone_name, epochs=args.epochs))
+    print(train(config, backbone_name, epochs=args.epochs, resume=args.resume))
 
 
 if __name__ == "__main__":